zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-29 14:58:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

shapeless slice update and broadcast when possible (#1727)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-23 11:25:15 -08:00
committed by GitHub
parent 0308e9af71
commit ebfe64b92d
6 changed files with 43 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/src/indexing.cpp
View File

@@ -766,7 +766,8 @@ auto mlx_slice_update(
const ScalarOrArray& v) {
// Can't route to slice update if not slice or tuple
if (src.ndim() == 0 ||
(!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj))) {
(!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj) &&
!nb::isinstance<nb::int_>(obj))) {
return std::make_pair(false, src);
}
if (nb::isinstance<nb::tuple>(obj)) {
@@ -777,7 +778,6 @@ auto mlx_slice_update(
}
}
}
// Should be able to route to slice update
// Pre process tuple
@@ -797,6 +797,20 @@ auto mlx_slice_update(
mx::Shape starts(src.ndim(), 0);
mx::Shape stops = src.shape();
mx::Shape strides(src.ndim(), 1);
if (nb::isinstance<nb::int_>(obj)) {
if (src.ndim() < 1) {
std::ostringstream msg;
msg << "Too many indices for array with " << src.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
auto idx = nb::cast<int>(obj);
idx = idx < 0 ? idx + stops[0] : idx;
starts[0] = idx;
stops[0] = idx + 1;
auto out = slice_update(
src, up, std::move(starts), std::move(stops), std::move(strides));
return std::make_pair(true, out);
}
// If it's just a simple slice, just do a slice update and return
if (nb::isinstance<nb::slice>(obj)) {

13
python/tests/test_compile.py
View File

@@ -817,6 +817,19 @@ class TestCompile(mlx_tests.MLXTestCase):
fun = mx.compile(lambda a, b: a @ b, shapeless=True)
self.assertTrue(mx.allclose(fun(a, b), a @ b))
def test_shapeless_compile_slice_update(self):
def fun(x):
x[2] = mx.array([3.0])
return x
cfun = mx.compile(fun, shapeless=True)
a = mx.array([0.0, 1.0, 2.0, 3.0])
self.assertTrue(mx.allclose(cfun(a), fun(a)))
a = mx.array([0.0, 1.0, 2.0, 3.0, 4.0])
self.assertTrue(mx.allclose(cfun(a), fun(a)))
if __name__ == "__main__":
unittest.main()