Multi output primitives (#330)

* Multi-output primitives

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

python/src/ops.cpp

@@ -303,6 +303,33 @@ void init_ops(py::module_& m) {
         Returns:
             array: The quotient ``a / b``.
       )pbdoc");
+  m.def(
+      "divmod",
+      [](const ScalarOrArray& a_, const ScalarOrArray& b_, StreamOrDevice s) {
+        auto [a, b] = to_arrays(a_, b_);
+        return divmod(a, b, s);
+      },
+      "a"_a,
+      "b"_a,
+      py::pos_only(),
+      py::kw_only(),
+      "stream"_a = none,
+      R"pbdoc(
+        divmod(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array
+
+        Element-wise quotient and remainder.
+
+        The fuction ``divmod(a, b)`` is equivalent to but faster than
+        ``(a // b, a % b)``. The function uses numpy-style broadcasting
+        semantics. Either or both input arrays can also be scalars.
+
+        Args:
+            a (array): Input array or scalar.
+            b (array): Input array or scalar.
+
+        Returns:
+            tuple(array, array): The quotient ``a // b`` and remainder ``a % b``.
+      )pbdoc");
   m.def(
       "floor_divide",
       [](const ScalarOrArray& a_, const ScalarOrArray& b_, StreamOrDevice s) {

python/tests/test_graph.py

@@ -0,0 +1,37 @@
+# Copyright © 2023 Apple Inc.
+
+import io
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+class TestGraph(mlx_tests.MLXTestCase):
+    def test_to_dot(self):
+        # Simply test that a few cases run.
+        # Nothing too specific about the graph format
+        # for now to keep it flexible
+        a = mx.array(1.0)
+        f = io.StringIO()
+        mx.export_to_dot(f, a)
+        f.seek(0)
+        self.assertTrue(len(f.read()) > 0)
+
+        b = mx.array(2.0)
+        c = a + b
+        f = io.StringIO()
+        mx.export_to_dot(f, c)
+        f.seek(0)
+        self.assertTrue(len(f.read()) > 0)
+
+        # Multi output case
+        c = mx.divmod(a, b)
+        f = io.StringIO()
+        mx.export_to_dot(f, *c)
+        f.seek(0)
+        self.assertTrue(len(f.read()) > 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/tests/test_ops.py

@@ -1314,7 +1314,7 @@ class TestOps(mlx_tests.MLXTestCase):
             for axis in (None, 0, 1, 2):
                 c_npy = npop(a_npy, axis=axis)
                 c_mlx = mxop(a_mlx, axis=axis)
-                self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-4, atol=1e-4))
+                self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3))
 
         for op in ["cumsum", "cumprod", "cummax", "cummin"]:
             c1 = mxop(a_mlx, axis=2)
@@ -1597,6 +1597,28 @@ class TestOps(mlx_tests.MLXTestCase):
             np.outer,
         )
 
+    def test_divmod(self):
+        # A few sizes for the inputs with and without broadcasting
+        sizes = [
+            ((1,), (1,)),
+            ((1,), (10,)),
+            ((10,), (1,)),
+            ((3,), (3,)),
+            ((2, 2, 2), (1, 2, 1)),
+            ((2, 1, 2), (1, 2, 1)),
+            ((2, 2, 2, 2), (2, 2, 2, 2)),
+        ]
+        types = [np.uint16, np.uint32, np.int32, np.float16, np.float32]
+        for s1, s2 in sizes:
+            for t in types:
+                a_np = np.random.uniform(1, 100, size=s1).astype(t)
+                b_np = np.random.uniform(1, 100, size=s2).astype(t)
+                np_out = np.divmod(a_np, b_np)
+                mx_out = mx.divmod(mx.array(a_np), mx.array(b_np))
+                self.assertTrue(
+                    np.allclose(np_out[0], mx_out[0]), msg=f"Shapes {s1} {s2}, Type {t}"
+                )
+
 
 if __name__ == "__main__":
     unittest.main()