From 1d018a0d790c94df2f3c254347ca5a864efb1099 Mon Sep 17 00:00:00 2001
From: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com>
Date: Fri, 19 Jun 2026 20:26:28 -0400
Subject: [PATCH 1/4] feat(array-api): add positive, logical_xor, trunc,
 count_nonzero, diff, full_like

New elementwise and utility ops required by the Python array API standard
(https://data-apis.org/array-api/latest/):

- positive: unary plus, returns a copy (mx::astype to same dtype)
- logical_xor: element-wise XOR via not_equal(bool(a), bool(b))
- trunc: truncate toward zero (where(a < 0, ceil, floor))
- count_nonzero: count non-zero elements; returns int32; supports axis/keepdims
- diff: n-th discrete difference along an axis, with optional prepend/append
- full_like: fill an array shaped like the input; optional dtype override

Docs and tests included.

Part of the array API split from #3684.
---
 docs/src/python/ops.rst  |   5 ++
 python/src/ops.cpp       | 158 +++++++++++++++++++++++++++++++++++++++
 python/tests/test_ops.py |  47 ++++++++++++
 3 files changed, 210 insertions(+)

diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst
index 84e0b9d08b..eb4dffc262 100644
--- a/docs/src/python/ops.rst
+++ b/docs/src/python/ops.rst
@@ -70,6 +70,7 @@ Operations
    dequantize
    diag
    diagonal
+   diff
    divide
    divmod
    einsum
@@ -87,6 +88,7 @@ Operations
    floor_divide
    full
    from_dlpack
+   full_like
    from_fp8
    gather_mm
    gather_qmm
@@ -121,6 +123,7 @@ Operations
    logical_not
    logical_and
    logical_or
+   logical_xor
    logsumexp
    matmul
    max
@@ -141,6 +144,7 @@ Operations
    partition
    pad
    permute_dims
+   positive
    power
    prod
    put_along_axis
@@ -195,6 +199,7 @@ Operations
    tri
    tril
    triu
+   trunc
    unflatten
    unstack
    vecdot
diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index f11f98427d..a85a6a3f14 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -5918,4 +5918,162 @@ void init_ops(nb::module_& m) {
   m.attr("empty_like") = m.attr("zeros_like");
   m.attr("matrix_transpose") = m.attr("transpose");
   m.attr("pow") = m.attr("power");
+  // Array API elementwise functions.
+  m.def(
+      "positive",
+      &mx::positive,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def positive(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise unary plus. Returns a copy of the input.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: A copy of ``a``.
+      )pbdoc");
+  m.def(
+      "logical_xor",
+      [](const ScalarOrArray& a_,
+         const ScalarOrArray& b_,
+         mx::StreamOrDevice s) {
+        auto [a, b] = to_arrays(a_, b_);
+        return mx::logical_xor(a, b, s);
+      },
+      nb::arg(),
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def logical_xor(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise logical exclusive or.
+
+        Args:
+            a (array): First input array or scalar.
+            b (array): Second input array or scalar.
+
+        Returns:
+            array: The boolean array containing the logical xor of ``a`` and ``b``.
+      )pbdoc");
+  m.def(
+      "trunc",
+      &mx::trunc,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def trunc(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise truncation towards zero.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: The truncated array.
+      )pbdoc");
+  m.def(
+      "count_nonzero",
+      [](const mx::array& a,
+         const IntOrVec& axis,
+         bool keepdims,
+         mx::StreamOrDevice s) {
+        if (std::holds_alternative<std::monostate>(axis)) {
+          return mx::count_nonzero(a, keepdims, s);
+        } else if (auto pv = std::get_if<int>(&axis); pv) {
+          return mx::count_nonzero(a, *pv, keepdims, s);
+        } else {
+          return mx::count_nonzero(
+              a, std::get<std::vector<int>>(axis), keepdims, s);
+        }
+      },
+      nb::arg(),
+      "axis"_a = nb::none(),
+      nb::kw_only(),
+      "keepdims"_a = false,
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def count_nonzero(a: array, /, *, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Count the number of non-zero elements along the given axis.
+
+        Args:
+            a (array): Input array.
+            axis (int or tuple(int), optional): Axis or axes to count over.
+              Defaults to ``None`` in which case the whole array is counted.
+            keepdims (bool, optional): Keep the reduced axes as size one.
+              Default: ``False``.
+
+        Returns:
+            array: The counts as an ``int32`` array.
+      )pbdoc");
+  m.def(
+      "diff",
+      [](const mx::array& a,
+         int n,
+         int axis,
+         const std::optional<mx::array>& prepend,
+         const std::optional<mx::array>& append,
+         mx::StreamOrDevice s) {
+        return mx::diff(a, n, axis, prepend, append, s);
+      },
+      nb::arg(),
+      "n"_a = 1,
+      "axis"_a = -1,
+      nb::kw_only(),
+      "prepend"_a = nb::none(),
+      "append"_a = nb::none(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def diff(a: array, /, n: int = 1, axis: int = -1, *, prepend: Optional[array] = None, append: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        The n-th discrete difference along the given axis.
+
+        Args:
+            a (array): Input array.
+            n (int, optional): The number of times to difference. Default: ``1``.
+            axis (int, optional): The axis along which to difference.
+              Default: ``-1``.
+            prepend (array, optional): Values to prepend along ``axis`` before
+              differencing.
+            append (array, optional): Values to append along ``axis`` before
+              differencing.
+
+        Returns:
+            array: The n-th differences.
+      )pbdoc");
+  // Array API creation functions.
+  m.def(
+      "full_like",
+      [](const mx::array& a,
+         const ScalarOrArray& vals,
+         std::optional<mx::Dtype> dtype,
+         mx::StreamOrDevice s) {
+        auto t = dtype.value_or(a.dtype());
+        return mx::full_like(a, to_array(vals, t), t, s);
+      },
+      nb::arg(),
+      "vals"_a,
+      "dtype"_a = nb::none(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def full_like(a: array, vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        An array filled with ``vals`` with the same shape as the input.
+
+        Args:
+            a (array): The input to take the shape from.
+            vals (float or int or array): Values to fill the array with.
+            dtype (Dtype, optional): Data type of the output array. If
+              unspecified the type of the input is used.
+
+        Returns:
+            array: The output array.
+      )pbdoc");
 }
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index 3220db7c79..4335d05a61 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -3481,6 +3481,53 @@ def test_to_from_fp8(self):
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
 
+    def test_array_api_elementwise(self):
+        a = mx.array([-1.5, -0.5, 0.0, 0.5, 2.7])
+        self.assertEqual(mx.positive(a).tolist(), a.tolist())
+        self.assertEqual(mx.trunc(a).tolist(), [-1.0, 0.0, 0.0, 0.0, 2.0])
+
+        x = mx.array([True, True, False, False])
+        y = mx.array([True, False, True, False])
+        self.assertEqual(mx.logical_xor(x, y).tolist(), [False, True, True, False])
+
+        c = mx.array([[0, 1, 0], [2, 3, 0]])
+        self.assertEqual(mx.count_nonzero(c).item(), 3)
+        self.assertEqual(mx.count_nonzero(c, axis=0).tolist(), [1, 2, 0])
+        self.assertEqual(mx.count_nonzero(c, axis=1).tolist(), [1, 2])
+        self.assertEqual(mx.count_nonzero(c).dtype, mx.int32)
+
+    def test_diff(self):
+        a = mx.array([1, 2, 4, 7, 0])
+        self.assertEqual(mx.diff(a).tolist(), [1, 2, 3, -7])
+        self.assertEqual(mx.diff(a, n=2).tolist(), [1, 1, -10])
+        self.assertEqual(mx.diff(a, n=0).tolist(), a.tolist())
+
+        m = mx.array([[1, 3, 6], [0, 5, 6]])
+        self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]])
+        self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]])
+
+        # prepend / append.
+        self.assertEqual(
+            mx.diff(mx.array([2, 4, 7]), prepend=mx.array([0])).tolist(),
+            [2, 2, 3],
+        )
+        self.assertEqual(
+            mx.diff(mx.array([2, 4, 7]), append=mx.array([10])).tolist(),
+            [2, 3, 3],
+        )
+
+        with self.assertRaises(ValueError):
+            mx.diff(a, axis=1)
+
+    def test_array_api_creation(self):
+        a = mx.arange(6, dtype=mx.int16).reshape(2, 3)
+
+        fl = mx.full_like(a, 7)
+        self.assertEqual(fl.shape, (2, 3))
+        self.assertEqual(fl.dtype, mx.int16)
+        self.assertTrue(mx.all(fl == 7).item())
+        self.assertEqual(mx.full_like(a, 1.5, dtype=mx.float32).dtype, mx.float32)
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

From 505d5563cc66e87e491a804ab37351dd923e1a4f Mon Sep 17 00:00:00 2001
From: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com>
Date: Sun, 21 Jun 2026 07:10:46 -0400
Subject: [PATCH 2/4] fix(array-api): move new ops to C++ (mlx/ops.h +
 mlx/ops.cpp)

Per zcbenz review: ops that are not pure aliases must have C++
implementations exposed through the usual mlx/ops.h + mlx/ops.cpp
path, not inline Python-binding lambdas.

- positive: array copy (like __copy__)
- logical_xor: not_equal(astype(a,bool_), astype(b,bool_))
- trunc: where(less(a,0), ceil, floor); integers returned as-is
- count_nonzero: sum(not_equal(a,0).astype(int32)) with axis overloads
- diff: n-th slice subtraction with optional prepend/append
- full_like: delegates to existing mx::full_like C++ overload

Python bindings in python/src/ops.cpp now call the C++ functions
instead of duplicating the logic inline.
---
 mlx/ops.cpp | 93 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 mlx/ops.h   | 36 +++++++++++++++++++++
 2 files changed, 129 insertions(+)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index d56ed7ffa4..e49db9567a 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -2158,6 +2158,32 @@ array sum(
   return sum(a, std::vector<int>{axis}, keepdims, s);
 }
 
+array count_nonzero(
+    const array& a,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return count_nonzero(a, axes, keepdims, s);
+}
+
+array count_nonzero(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return count_nonzero(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array count_nonzero(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  auto nz = astype(not_equal(a, array(0, a.dtype()), s), int32, s);
+  return sum(nz, axes, keepdims, s);
+}
+
 array mean(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
   std::vector<int> axes(a.ndim());
   std::iota(axes.begin(), axes.end(), 0);
@@ -2848,6 +2874,10 @@ array abs(const array& a, StreamOrDevice s /* = {} */) {
   return out;
 }
 
+array positive(const array& a, StreamOrDevice s /* = {} */) {
+  return array(a);
+}
+
 array negative(const array& a, StreamOrDevice s /* = {} */) {
   if (a.dtype() == bool_) {
     auto msg = "[negative] Not supported for bool, use logical_not instead.";
@@ -2900,6 +2930,10 @@ array operator||(const array& a, const array& b) {
   return logical_or(a, b);
 }
 
+array logical_xor(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  return not_equal(astype(a, bool_, s), astype(b, bool_, s), s);
+}
+
 array reciprocal(const array& a, StreamOrDevice s /* = {} */) {
   auto dtype = at_least_float(a.dtype());
   return divide(array(1.0f, dtype), a, to_stream(s));
@@ -3052,6 +3086,17 @@ array ceil(const array& a, StreamOrDevice s /* = {} */) {
   return array(a.shape(), a.dtype(), std::make_shared<Ceil>(to_stream(s)), {a});
 }
 
+array trunc(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == complex64) {
+    throw std::invalid_argument("[trunc] Not supported for complex64.");
+  }
+  if (issubdtype(a.dtype(), integer)) {
+    return array(a);
+  }
+  auto zero = array(0, a.dtype());
+  return where(less(a, zero, s), ceil(a, s), floor(a, s), s);
+}
+
 array square(const array& a, StreamOrDevice s /* = {} */) {
   return array(
       a.shape(), a.dtype(), std::make_shared<Square>(to_stream(s)), {a});
@@ -4063,6 +4108,54 @@ array cummin(
   return cummin(flatten(a, s), 0, reverse, inclusive, s);
 }
 
+array diff(
+    const array& a,
+    int n /* = 1 */,
+    int axis /* = -1 */,
+    StreamOrDevice s /* = {} */) {
+  return diff(a, n, axis, std::nullopt, std::nullopt, s);
+}
+
+array diff(
+    const array& a,
+    int n,
+    int axis,
+    const std::optional<array>& prepend,
+    const std::optional<array>& append,
+    StreamOrDevice s /* = {} */) {
+  int ndim = static_cast<int>(a.ndim());
+  int ax = axis < 0 ? axis + ndim : axis;
+  if (ax < 0 || ax >= ndim) {
+    throw std::invalid_argument("[diff] Axis is out of bounds for the array.");
+  }
+  if (n < 0) {
+    throw std::invalid_argument("[diff] Order `n` must be non-negative.");
+  }
+  array x = a;
+  if (prepend || append) {
+    std::vector<array> parts;
+    if (prepend) {
+      parts.push_back(*prepend);
+    }
+    parts.push_back(x);
+    if (append) {
+      parts.push_back(*append);
+    }
+    x = concatenate(parts, ax, s);
+  }
+  for (int i = 0; i < n; ++i) {
+    Shape upper_start(x.ndim(), 0);
+    Shape lower_stop = x.shape();
+    Shape strides(x.ndim(), 1);
+    upper_start[ax] = 1;
+    lower_stop[ax] = x.shape(ax) - 1;
+    auto upper = slice(x, upper_start, x.shape(), strides, s);
+    auto lower = slice(x, Shape(x.ndim(), 0), lower_stop, strides, s);
+    x = subtract(upper, lower, s);
+  }
+  return x;
+}
+
 array logcumsumexp(
     const array& a,
     int axis,
diff --git a/mlx/ops.h b/mlx/ops.h
index 97f06eb6e3..1d2c8544ea 100644
--- a/mlx/ops.h
+++ b/mlx/ops.h
@@ -622,6 +622,20 @@ sum(const array& a,
 MLX_API array
 sum(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
 
+/** Count the number of non-zero elements in an array. */
+MLX_API array
+count_nonzero(const array& a, bool keepdims = false, StreamOrDevice s = {});
+MLX_API array count_nonzero(
+    const array& a,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+MLX_API array count_nonzero(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
 /** Computes the mean of the elements of an array. */
 MLX_API array mean(const array& a, bool keepdims, StreamOrDevice s = {});
 inline array mean(const array& a, StreamOrDevice s = {}) {
@@ -883,6 +897,9 @@ MLX_API array logsumexp(
 /** Absolute value of elements in an array. */
 MLX_API array abs(const array& a, StreamOrDevice s = {});
 
+/** Unary plus — return a copy of the array unchanged. */
+MLX_API array positive(const array& a, StreamOrDevice s = {});
+
 /** Negate an array. */
 MLX_API array negative(const array& a, StreamOrDevice s = {});
 MLX_API array operator-(const array& a);
@@ -902,6 +919,10 @@ MLX_API array operator&&(const array& a, const array& b);
 MLX_API array logical_or(const array& a, const array& b, StreamOrDevice s = {});
 MLX_API array operator||(const array& a, const array& b);
 
+/** Logical exclusive or of two arrays */
+MLX_API array
+logical_xor(const array& a, const array& b, StreamOrDevice s = {});
+
 /** The reciprocal (1/x) of the elements in an array. */
 MLX_API array reciprocal(const array& a, StreamOrDevice s = {});
 
@@ -979,6 +1000,9 @@ MLX_API array floor(const array& a, StreamOrDevice s = {});
 /** Ceil the element of an array. **/
 MLX_API array ceil(const array& a, StreamOrDevice s = {});
 
+/** Truncate the elements of an array towards zero. **/
+MLX_API array trunc(const array& a, StreamOrDevice s = {});
+
 /** Square the elements of an array. */
 MLX_API array square(const array& a, StreamOrDevice s = {});
 
@@ -1388,6 +1412,18 @@ MLX_API array cummin(
     bool inclusive = true,
     StreamOrDevice s = {});
 
+/** The n-th discrete difference along the given axis. */
+MLX_API array
+diff(const array& a, int n = 1, int axis = -1, StreamOrDevice s = {});
+
+MLX_API array diff(
+    const array& a,
+    int n,
+    int axis,
+    const std::optional<array>& prepend,
+    const std::optional<array>& append,
+    StreamOrDevice s = {});
+
 /** General convolution with a filter */
 MLX_API array conv_general(
     array input,

From 80bf4b7f1342cdb81baf9cad3e9cdabc30427ab5 Mon Sep 17 00:00:00 2001
From: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com>
Date: Fri, 26 Jun 2026 07:31:10 -0400
Subject: [PATCH 3/4] Address review: drop prepend/append args from diff

Per maintainer request (#3730), remove the prepend/append parameters and
the second diff() overload. diff is now diff(a, n, axis, stream) only;
callers that need edge padding can concatenate before calling.
---
 mlx/ops.cpp              | 21 ---------------------
 mlx/ops.h                |  8 --------
 python/src/ops.cpp       | 17 +++--------------
 python/tests/test_ops.py | 10 ----------
 4 files changed, 3 insertions(+), 53 deletions(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index e49db9567a..a26d8f25d2 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -4113,16 +4113,6 @@ array diff(
     int n /* = 1 */,
     int axis /* = -1 */,
     StreamOrDevice s /* = {} */) {
-  return diff(a, n, axis, std::nullopt, std::nullopt, s);
-}
-
-array diff(
-    const array& a,
-    int n,
-    int axis,
-    const std::optional<array>& prepend,
-    const std::optional<array>& append,
-    StreamOrDevice s /* = {} */) {
   int ndim = static_cast<int>(a.ndim());
   int ax = axis < 0 ? axis + ndim : axis;
   if (ax < 0 || ax >= ndim) {
@@ -4132,17 +4122,6 @@ array diff(
     throw std::invalid_argument("[diff] Order `n` must be non-negative.");
   }
   array x = a;
-  if (prepend || append) {
-    std::vector<array> parts;
-    if (prepend) {
-      parts.push_back(*prepend);
-    }
-    parts.push_back(x);
-    if (append) {
-      parts.push_back(*append);
-    }
-    x = concatenate(parts, ax, s);
-  }
   for (int i = 0; i < n; ++i) {
     Shape upper_start(x.ndim(), 0);
     Shape lower_stop = x.shape();
diff --git a/mlx/ops.h b/mlx/ops.h
index 1d2c8544ea..40c8d404b1 100644
--- a/mlx/ops.h
+++ b/mlx/ops.h
@@ -1416,14 +1416,6 @@ MLX_API array cummin(
 MLX_API array
 diff(const array& a, int n = 1, int axis = -1, StreamOrDevice s = {});
 
-MLX_API array diff(
-    const array& a,
-    int n,
-    int axis,
-    const std::optional<array>& prepend,
-    const std::optional<array>& append,
-    StreamOrDevice s = {});
-
 /** General convolution with a filter */
 MLX_API array conv_general(
     array input,
diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index a85a6a3f14..5ebff106f5 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -6014,23 +6014,16 @@ void init_ops(nb::module_& m) {
       )pbdoc");
   m.def(
       "diff",
-      [](const mx::array& a,
-         int n,
-         int axis,
-         const std::optional<mx::array>& prepend,
-         const std::optional<mx::array>& append,
-         mx::StreamOrDevice s) {
-        return mx::diff(a, n, axis, prepend, append, s);
+      [](const mx::array& a, int n, int axis, mx::StreamOrDevice s) {
+        return mx::diff(a, n, axis, s);
       },
       nb::arg(),
       "n"_a = 1,
       "axis"_a = -1,
       nb::kw_only(),
-      "prepend"_a = nb::none(),
-      "append"_a = nb::none(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def diff(a: array, /, n: int = 1, axis: int = -1, *, prepend: Optional[array] = None, append: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"),
+          "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
         The n-th discrete difference along the given axis.
 
@@ -6039,10 +6032,6 @@ void init_ops(nb::module_& m) {
             n (int, optional): The number of times to difference. Default: ``1``.
             axis (int, optional): The axis along which to difference.
               Default: ``-1``.
-            prepend (array, optional): Values to prepend along ``axis`` before
-              differencing.
-            append (array, optional): Values to append along ``axis`` before
-              differencing.
 
         Returns:
             array: The n-th differences.
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index 4335d05a61..cb3db22c24 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -3506,16 +3506,6 @@ def test_diff(self):
         self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]])
         self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]])
 
-        # prepend / append.
-        self.assertEqual(
-            mx.diff(mx.array([2, 4, 7]), prepend=mx.array([0])).tolist(),
-            [2, 2, 3],
-        )
-        self.assertEqual(
-            mx.diff(mx.array([2, 4, 7]), append=mx.array([10])).tolist(),
-            [2, 3, 3],
-        )
-
         with self.assertRaises(ValueError):
             mx.diff(a, axis=1)
 

From a1151b56a5e9bcdec73367757c257820a531e05b Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Sun, 28 Jun 2026 13:28:47 +0900
Subject: [PATCH 4/4] Rearrange code

---
 docs/src/python/ops.rst  |   1 +
 python/src/ops.cpp       | 287 +++++++++++++++++++--------------------
 python/tests/test_ops.py |  66 ++++-----
 3 files changed, 170 insertions(+), 184 deletions(-)

diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst
index eb4dffc262..2303ed80c6 100644
--- a/docs/src/python/ops.rst
+++ b/docs/src/python/ops.rst
@@ -65,6 +65,7 @@ Operations
    cummin
    cumprod
    cumsum
+   count_nonzero
    degrees
    depends
    dequantize
diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index 5ebff106f5..c709ec86b1 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -297,6 +297,23 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The sign of ``a``.
       )pbdoc");
+  m.def(
+      "positive",
+      &mx::positive,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def positive(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise unary plus. Returns a copy of the input.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: A copy of ``a``.
+      )pbdoc");
   m.def(
       "negative",
       [](const ScalarOrArray& a, mx::StreamOrDevice s) {
@@ -733,6 +750,23 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The matrix product of ``a`` and ``b``.
       )pbdoc");
+  m.def(
+      "trunc",
+      &mx::trunc,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def trunc(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise truncation towards zero.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: The truncated array.
+      )pbdoc");
   m.def(
       "square",
       [](const ScalarOrArray& a, mx::StreamOrDevice s) {
@@ -871,6 +905,27 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The boolean array containing the logical or of ``a`` and ``b``.
     )pbdoc");
+  m.def(
+      "logical_xor",
+      [](const ScalarOrArray& a, const ScalarOrArray& b, mx::StreamOrDevice s) {
+        return mx::logical_xor(to_array(a), to_array(b), s);
+      },
+      nb::arg(),
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def logical_xor(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise logical exclusive or.
+
+        Args:
+            a (array): First input array or scalar.
+            b (array): Second input array or scalar.
+
+        Returns:
+            array: The boolean array containing the logical xor of ``a`` and ``b``.
+      )pbdoc");
   m.def(
       "logaddexp",
       [](const ScalarOrArray& a_,
@@ -1792,6 +1847,34 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The output array with the specified shape and values.
       )pbdoc");
+  m.def(
+      "full_like",
+      [](const mx::array& a,
+         const ScalarOrArray& vals,
+         std::optional<mx::Dtype> dtype,
+         mx::StreamOrDevice s) {
+        auto t = dtype.value_or(a.dtype());
+        return mx::full_like(a, to_array(vals, t), t, s);
+      },
+      nb::arg(),
+      "vals"_a,
+      "dtype"_a = nb::none(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def full_like(a: array, vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        An array filled with ``vals`` with the same shape as the input.
+
+        Args:
+            a (array): The input to take the shape from.
+            vals (float or int or array): Values to fill the array with.
+            dtype (Dtype, optional): Data type of the output array. If
+              unspecified the type of the input is used.
+
+        Returns:
+            array: The output array.
+      )pbdoc");
   m.def(
       "zeros",
       [](const nb::object& shape,
@@ -2494,6 +2577,41 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The output array with the corresponding axes reduced.
       )pbdoc");
+  m.def(
+      "count_nonzero",
+      [](const mx::array& a,
+         const IntOrVec& axis,
+         bool keepdims,
+         mx::StreamOrDevice s) {
+        if (std::holds_alternative<std::monostate>(axis)) {
+          return mx::count_nonzero(a, keepdims, s);
+        } else if (auto pv = std::get_if<int>(&axis); pv) {
+          return mx::count_nonzero(a, *pv, keepdims, s);
+        } else {
+          return mx::count_nonzero(
+              a, std::get<std::vector<int>>(axis), keepdims, s);
+        }
+      },
+      nb::arg(),
+      "axis"_a = nb::none(),
+      nb::kw_only(),
+      "keepdims"_a = false,
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def count_nonzero(a: array, /, *, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Count the number of non-zero elements along the given axis.
+
+        Args:
+            a (array): Input array.
+            axis (int or tuple(int), optional): Axis or axes to count over.
+              Defaults to ``None`` in which case the whole array is counted.
+            keepdims (bool, optional): Keep the reduced axes as size one.
+              Default: ``False``.
+
+        Returns:
+            array: The counts as an ``int32`` array.
+      )pbdoc");
   m.def(
       "prod",
       [](const mx::array& a,
@@ -3585,6 +3703,28 @@ void init_ops(nb::module_& m) {
         Returns:
           array: The output array.
       )pbdoc");
+  m.def(
+      "diff",
+      &mx::diff,
+      nb::arg(),
+      "n"_a = 1,
+      "axis"_a = -1,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        The n-th discrete difference along the given axis.
+
+        Args:
+            a (array): Input array.
+            n (int, optional): The number of times to difference. Default: ``1``.
+            axis (int, optional): The axis along which to difference.
+              Default: ``-1``.
+
+        Returns:
+            array: The n-th differences.
+      )pbdoc");
   m.def(
       "conj",
       [](const ScalarOrArray& a, mx::StreamOrDevice s) {
@@ -5918,151 +6058,4 @@ void init_ops(nb::module_& m) {
   m.attr("empty_like") = m.attr("zeros_like");
   m.attr("matrix_transpose") = m.attr("transpose");
   m.attr("pow") = m.attr("power");
-  // Array API elementwise functions.
-  m.def(
-      "positive",
-      &mx::positive,
-      nb::arg(),
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def positive(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        Element-wise unary plus. Returns a copy of the input.
-
-        Args:
-            a (array): Input array.
-
-        Returns:
-            array: A copy of ``a``.
-      )pbdoc");
-  m.def(
-      "logical_xor",
-      [](const ScalarOrArray& a_,
-         const ScalarOrArray& b_,
-         mx::StreamOrDevice s) {
-        auto [a, b] = to_arrays(a_, b_);
-        return mx::logical_xor(a, b, s);
-      },
-      nb::arg(),
-      nb::arg(),
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def logical_xor(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        Element-wise logical exclusive or.
-
-        Args:
-            a (array): First input array or scalar.
-            b (array): Second input array or scalar.
-
-        Returns:
-            array: The boolean array containing the logical xor of ``a`` and ``b``.
-      )pbdoc");
-  m.def(
-      "trunc",
-      &mx::trunc,
-      nb::arg(),
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def trunc(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        Element-wise truncation towards zero.
-
-        Args:
-            a (array): Input array.
-
-        Returns:
-            array: The truncated array.
-      )pbdoc");
-  m.def(
-      "count_nonzero",
-      [](const mx::array& a,
-         const IntOrVec& axis,
-         bool keepdims,
-         mx::StreamOrDevice s) {
-        if (std::holds_alternative<std::monostate>(axis)) {
-          return mx::count_nonzero(a, keepdims, s);
-        } else if (auto pv = std::get_if<int>(&axis); pv) {
-          return mx::count_nonzero(a, *pv, keepdims, s);
-        } else {
-          return mx::count_nonzero(
-              a, std::get<std::vector<int>>(axis), keepdims, s);
-        }
-      },
-      nb::arg(),
-      "axis"_a = nb::none(),
-      nb::kw_only(),
-      "keepdims"_a = false,
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def count_nonzero(a: array, /, *, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        Count the number of non-zero elements along the given axis.
-
-        Args:
-            a (array): Input array.
-            axis (int or tuple(int), optional): Axis or axes to count over.
-              Defaults to ``None`` in which case the whole array is counted.
-            keepdims (bool, optional): Keep the reduced axes as size one.
-              Default: ``False``.
-
-        Returns:
-            array: The counts as an ``int32`` array.
-      )pbdoc");
-  m.def(
-      "diff",
-      [](const mx::array& a, int n, int axis, mx::StreamOrDevice s) {
-        return mx::diff(a, n, axis, s);
-      },
-      nb::arg(),
-      "n"_a = 1,
-      "axis"_a = -1,
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        The n-th discrete difference along the given axis.
-
-        Args:
-            a (array): Input array.
-            n (int, optional): The number of times to difference. Default: ``1``.
-            axis (int, optional): The axis along which to difference.
-              Default: ``-1``.
-
-        Returns:
-            array: The n-th differences.
-      )pbdoc");
-  // Array API creation functions.
-  m.def(
-      "full_like",
-      [](const mx::array& a,
-         const ScalarOrArray& vals,
-         std::optional<mx::Dtype> dtype,
-         mx::StreamOrDevice s) {
-        auto t = dtype.value_or(a.dtype());
-        return mx::full_like(a, to_array(vals, t), t, s);
-      },
-      nb::arg(),
-      "vals"_a,
-      "dtype"_a = nb::none(),
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      nb::sig(
-          "def full_like(a: array, vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
-      R"pbdoc(
-        An array filled with ``vals`` with the same shape as the input.
-
-        Args:
-            a (array): The input to take the shape from.
-            vals (float or int or array): Values to fill the array with.
-            dtype (Dtype, optional): Data type of the output array. If
-              unspecified the type of the input is used.
-
-        Returns:
-            array: The output array.
-      )pbdoc");
 }
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index cb3db22c24..8dd09217b1 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -702,6 +702,13 @@ def test_sum(self):
 
         self.assertTrue(np.array_equal(y_mlx, y_npy))
 
+    def test_count_nonzero(self):
+        c = mx.array([[0, 1, 0], [2, 3, 0]])
+        self.assertEqual(mx.count_nonzero(c).item(), 3)
+        self.assertEqual(mx.count_nonzero(c, axis=0).tolist(), [1, 2, 0])
+        self.assertEqual(mx.count_nonzero(c, axis=1).tolist(), [1, 2])
+        self.assertEqual(mx.count_nonzero(c).dtype, mx.int32)
+
     def test_prod(self):
         x = mx.array(
             [
@@ -939,6 +946,15 @@ def test_logical_or(self):
         result = a | b
         self.assertTrue(np.array_equal(result, expected))
 
+    def test_logical_xor(self):
+        x = mx.array([True, True, False, False])
+        y = mx.array([True, False, True, False])
+        self.assertEqual(mx.logical_xor(x, y).tolist(), [False, True, True, False])
+
+    def test_trunc(self):
+        a = mx.array([-1.5, -0.5, 0.0, 0.5, 2.7])
+        self.assertEqual(mx.trunc(a).tolist(), [-1.0, 0.0, 0.0, 0.0, 2.0])
+
     def test_square(self):
         a = mx.array([0.1, 0.5, 1.0, 10.0])
         result = mx.square(a)
@@ -2277,6 +2293,19 @@ def fn(its):
         mem4 = mx.get_peak_memory()
         self.assertEqual(mem2, mem4)
 
+    def test_diff(self):
+        a = mx.array([1, 2, 4, 7, 0])
+        self.assertEqual(mx.diff(a).tolist(), [1, 2, 3, -7])
+        self.assertEqual(mx.diff(a, n=2).tolist(), [1, 1, -10])
+        self.assertEqual(mx.diff(a, n=0).tolist(), a.tolist())
+
+        m = mx.array([[1, 3, 6], [0, 5, 6]])
+        self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]])
+        self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]])
+
+        with self.assertRaises(ValueError):
+            mx.diff(a, axis=1)
+
     def test_squeeze_expand(self):
         a = mx.zeros((2, 1, 2, 1))
         self.assertEqual(mx.squeeze(a).shape, (2, 2))
@@ -3481,43 +3510,6 @@ def test_to_from_fp8(self):
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
 
-    def test_array_api_elementwise(self):
-        a = mx.array([-1.5, -0.5, 0.0, 0.5, 2.7])
-        self.assertEqual(mx.positive(a).tolist(), a.tolist())
-        self.assertEqual(mx.trunc(a).tolist(), [-1.0, 0.0, 0.0, 0.0, 2.0])
-
-        x = mx.array([True, True, False, False])
-        y = mx.array([True, False, True, False])
-        self.assertEqual(mx.logical_xor(x, y).tolist(), [False, True, True, False])
-
-        c = mx.array([[0, 1, 0], [2, 3, 0]])
-        self.assertEqual(mx.count_nonzero(c).item(), 3)
-        self.assertEqual(mx.count_nonzero(c, axis=0).tolist(), [1, 2, 0])
-        self.assertEqual(mx.count_nonzero(c, axis=1).tolist(), [1, 2])
-        self.assertEqual(mx.count_nonzero(c).dtype, mx.int32)
-
-    def test_diff(self):
-        a = mx.array([1, 2, 4, 7, 0])
-        self.assertEqual(mx.diff(a).tolist(), [1, 2, 3, -7])
-        self.assertEqual(mx.diff(a, n=2).tolist(), [1, 1, -10])
-        self.assertEqual(mx.diff(a, n=0).tolist(), a.tolist())
-
-        m = mx.array([[1, 3, 6], [0, 5, 6]])
-        self.assertEqual(mx.diff(m, axis=0).tolist(), [[-1, 2, 0]])
-        self.assertEqual(mx.diff(m, axis=1).tolist(), [[2, 3], [5, 1]])
-
-        with self.assertRaises(ValueError):
-            mx.diff(a, axis=1)
-
-    def test_array_api_creation(self):
-        a = mx.arange(6, dtype=mx.int16).reshape(2, 3)
-
-        fl = mx.full_like(a, 7)
-        self.assertEqual(fl.shape, (2, 3))
-        self.assertEqual(fl.dtype, mx.int16)
-        self.assertTrue(mx.all(fl == 7).item())
-        self.assertEqual(mx.full_like(a, 1.5, dtype=mx.float32).dtype, mx.float32)
-
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()