zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-09 13:25:32 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

expose depends (#2606)

Browse Source
This commit is contained in:
Awni Hannun
2025-09-18 10:06:15 -07:00
committed by GitHub
parent 3f730e77aa
commit 50cc09887f
2 changed files with 53 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
python/src/ops.cpp
View File

@@ -5319,4 +5319,44 @@ void init_ops(nb::module_& m) {
>>> mx.broadcast_shapes((5, 1, 4), (1, 3, 1))
(5, 3, 4)
)pbdoc");
m.def(
"depends",
[](const nb::object& inputs_, const nb::object& deps_) {
bool return_vec = false;
std::vector<mx::array> inputs;
std::vector<mx::array> deps;
if (nb::isinstance<mx::array>(inputs_)) {
inputs = {nb::cast<mx::array>(inputs_)};
} else {
return_vec = true;
inputs = {nb::cast<std::vector<mx::array>>(inputs_)};
}
if (nb::isinstance<mx::array>(deps_)) {
deps = {nb::cast<mx::array>(deps_)};
} else {
deps = {nb::cast<std::vector<mx::array>>(deps_)};
}
auto out = depends(inputs, deps);
if (return_vec) {
return nb::cast(out);
} else {
return nb::cast(out[0]);
}
},
nb::arg(),
nb::arg(),
nb::sig(
"def depends(inputs: Union[array, Sequence[array]], dependencies: Union[array, Sequence[array]])"),
R"pbdoc(
Insert dependencies between arrays in the graph. The outputs are
identical to ``inputs`` but with dependencies on ``dependencies``.
Args:
inputs (array or Sequence[array]): The input array or arrays.
dependencies (array or Sequence[array]): The array or arrays
to insert dependencies on.
Returns:
array or Sequence[array]: The outputs which depend on dependencies.
)pbdoc");
}

13
python/tests/test_ops.py
View File

@@ -3081,6 +3081,19 @@ class TestOps(mlx_tests.MLXTestCase):
# Doesn't hang
x = mx.power(2, -1)
def test_depends(self):
a = mx.array([1.0, 2.0, 3.0])
b = mx.exp(a)
c = mx.log(a)
out = mx.depends([b], [c])[0]
self.assertTrue(mx.array_equal(out, b))
a = mx.array([1.0, 2.0, 3.0])
b = mx.exp(a)
c = mx.log(a)
out = mx.depends(b, c)
self.assertTrue(mx.array_equal(out, b))
class TestBroadcast(mlx_tests.MLXTestCase):
def test_broadcast_shapes(self):