add mx.trace (#1143) (#1147)

* working c++ trace implementation

* updated throw + added overloads

* added python binding for trace function

* pre-commit reformatting

* add trace to docs

* resolve comments

* remove to_stream call

python/src/ops.cpp

@@ -4065,6 +4065,45 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The extracted diagonal or the constructed diagonal matrix.
         )pbdoc");
+  m.def(
+      "trace",
+      [](const array& a,
+         int offset,
+         int axis1,
+         int axis2,
+         std::optional<Dtype> dtype,
+         StreamOrDevice s) {
+        if (!dtype.has_value()) {
+          return trace(a, offset, axis1, axis2, s);
+        }
+        return trace(a, offset, axis1, axis2, dtype.value(), s);
+      },
+      nb::arg(),
+      "offset"_a = 0,
+      "axis1"_a = 0,
+      "axis2"_a = 1,
+      "dtype"_a = nb::none(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype = Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Return the sum along a specified diagonal in the given array.
+
+        Args:
+          a (array): Input array
+          offset (int, optional): Offset of the diagonal from the main diagonal.
+            Can be positive or negative. Default: ``0``.
+          axis1 (int, optional): The first axis of the 2-D sub-arrays from which
+              the diagonals should be taken. Default: ``0``.
+          axis2 (int, optional): The second axis of the 2-D sub-arrays from which
+              the diagonals should be taken. Default: ``1``.
+          dtype (Dtype, optional): Data type of the output array. If
+              unspecified the output type is inferred from the input array.
+
+        Returns:
+            array: Sum of specified diagonal.
+        )pbdoc");
   m.def(
       "atleast_1d",
       [](const nb::args& arys, StreamOrDevice s) -> nb::object {

python/tests/test_ops.py

@@ -2091,6 +2091,38 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = mx.array(np.diag(x, k=-1))
         self.assertTrue(mx.array_equal(result, expected))
 
+    def test_trace(self):
+        a_mx = mx.arange(9, dtype=mx.int64).reshape((3, 3))
+        a_np = np.arange(9, dtype=np.int64).reshape((3, 3))
+
+        # Test 2D array
+        result = mx.trace(a_mx)
+        expected = np.trace(a_np)
+        self.assertEqualArray(result, mx.array(expected))
+
+        # Test dtype
+        result = mx.trace(a_mx, dtype=mx.float16)
+        expected = np.trace(a_np, dtype=np.float16)
+        self.assertEqualArray(result, mx.array(expected))
+
+        # Test offset
+        result = mx.trace(a_mx, offset=1)
+        expected = np.trace(a_np, offset=1)
+        self.assertEqualArray(result, mx.array(expected))
+
+        # Test axis1 and axis2
+        b_mx = mx.arange(27, dtype=mx.int64).reshape(3, 3, 3)
+        b_np = np.arange(27, dtype=np.int64).reshape(3, 3, 3)
+
+        result = mx.trace(b_mx, axis1=1, axis2=2)
+        expected = np.trace(b_np, axis1=1, axis2=2)
+        self.assertEqualArray(result, mx.array(expected))
+
+        # Test offset, axis1, axis2, and dtype
+        result = mx.trace(b_mx, offset=1, axis1=1, axis2=2, dtype=mx.float32)
+        expected = np.trace(b_np, offset=1, axis1=1, axis2=2, dtype=np.float32)
+        self.assertEqualArray(result, mx.array(expected))
+
     def test_atleast_1d(self):
         def compare_nested_lists(x, y):
             if isinstance(x, list) and isinstance(y, list):