zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Support more Numpy interfaces for masked_scatter (#2832)

Browse Source
This commit is contained in:
CCYeh
2025-12-02 02:51:02 +01:00
committed by GitHub
parent 6e762fe2e2
commit 8879ee00eb
4 changed files with 34 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/src/indexing.cpp
View File

@@ -766,7 +766,7 @@ auto mlx_slice_update(
const nb::object& obj,
const ScalarOrArray& v) {
// Can't route to slice update if not slice, tuple, or int
if (src.ndim() == 0 ||
if (src.ndim() == 0 || nb::isinstance<nb::bool_>(obj) ||
(!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj) &&
!nb::isinstance<nb::int_>(obj))) {
return std::make_pair(false, src);
@@ -888,7 +888,9 @@ auto mlx_slice_update(
std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
using NDArray = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>;
if (nb::isinstance<mx::array>(obj)) {
if (nb::isinstance<nb::bool_>(obj)) {
return mx::array(nb::cast<bool>(obj), mx::bool_);
} else if (nb::isinstance<mx::array>(obj)) {
auto mask = nb::cast<mx::array>(obj);
if (mask.dtype() == mx::bool_) {
return mask;
@@ -898,6 +900,11 @@ std::optional<mx::array> extract_boolean_mask(const nb::object& obj) {
if (mask.dtype() == nb::dtype<bool>()) {
return nd_array_to_mlx(mask, mx::bool_);
}
} else if (nb::isinstance<nb::list>(obj)) {
auto mask = array_from_list(nb::cast<nb::list>(obj), {});
if (mask.dtype() == mx::bool_) {
return mask;
}
}
return std::nullopt;
}

23
python/tests/test_array.py
View File

@@ -1929,8 +1929,27 @@ class TestArray(mlx_tests.MLXTestCase):
self.assertTrue(np.array_equal(a, anp))
def test_setitem_with_boolean_mask(self):
mask_np = np.zeros((10, 10), dtype=bool)
mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0
# Python list mask
a = mx.array([1.0, 2.0, 3.0])
mask = [True, False, True]
src = mx.array([5.0, 6.0])
expected = mx.array([5.0, 2.0, 6.0])
a[mask] = src
self.assertTrue(mx.array_equal(a, expected))
# mx.array scalar mask
a = mx.array([1.0, 2.0, 3.0])
mask = mx.array(True)
expected = mx.array([5.0, 5.0, 5.0])
a[mask] = 5.0
self.assertTrue(mx.array_equal(a, expected))
# scalar mask
a = mx.array([1.0, 2.0, 3.0])
mask = True
expected = mx.array([5.0, 5.0, 5.0])
a[mask] = 5.0
self.assertTrue(mx.array_equal(a, expected))
mask_np = np.zeros((1, 10, 10), dtype=bool)
with self.assertRaises(ValueError):