diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 84e0b9d08b..2303ed80c6 100644 --- a/docs/src/python/ops.rst +++ b/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 diff --git a/mlx/ops.cpp b/mlx/ops.cpp index d56ed7ffa4..a26d8f25d2 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2158,6 +2158,32 @@ array sum( return sum(a, std::vector{axis}, keepdims, s); } +array count_nonzero( + const array& a, + bool keepdims /* = false */, + StreamOrDevice s /* = {} */) { + std::vector 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{axis}, keepdims, s); +} + +array count_nonzero( + const array& a, + const std::vector& 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 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(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(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(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, diff --git a/mlx/ops.h b/mlx/ops.h index 97f06eb6e3..40c8d404b1 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& 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, diff --git a/python/src/ops.cpp b/python/src/ops.cpp index f11f98427d..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 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(axis)) { + return mx::count_nonzero(a, keepdims, s); + } else if (auto pv = std::get_if(&axis); pv) { + return mx::count_nonzero(a, *pv, keepdims, s); + } else { + return mx::count_nonzero( + a, std::get>(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) { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 3220db7c79..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))