array API: add positive, logical_xor, trunc, count_nonzero, diff, full_like (#3730)

Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com>
Co-authored-by: Cheng <git@zcbenz.com>

Author: katlun-lgtm

docs/src/python/ops.rst

@@ -65,11 +65,13 @@ Operations
    cummin
    cumprod
    cumsum
+   count_nonzero
    degrees
    depends
    dequantize
    diag
    diagonal
+   diff
    divide
    divmod
    einsum
@@ -87,6 +89,7 @@ Operations
    floor_divide
    full
    from_dlpack
+   full_like
    from_fp8
    gather_mm
    gather_qmm
@@ -121,6 +124,7 @@ Operations
    logical_not
    logical_and
    logical_or
+   logical_xor
    logsumexp
    matmul
    max
@@ -141,6 +145,7 @@ Operations
    partition
    pad
    permute_dims
+   positive
    power
    prod
    put_along_axis
@@ -195,6 +200,7 @@ Operations
    tri
    tril
    triu
+   trunc
    unflatten
    unstack
    vecdot

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,33 @@ array cummin(
   return cummin(flatten(a, s), 0, reverse, inclusive, s);
 }
 
+array diff(
+    const array& a,
+    int n /* = 1 */,
+    int axis /* = -1 */,
+    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;
+  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,

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,10 @@ 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 = {});
+
 /** General convolution with a filter */
 MLX_API array conv_general(
     array input,

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) {