diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index de40c548a2..bdf6d538b5 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -80,6 +80,7 @@ Operations expand_dims eye flatten + flip floor floor_divide full @@ -190,6 +191,7 @@ Operations tril triu unflatten + unstack var view where diff --git a/mlx/ops.cpp b/mlx/ops.cpp index e4ce3d750f..9978d148da 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -645,6 +645,33 @@ array expand_dims( return expand_dims_impl(a, std::move(sorted_axes), s); } +array flip( + const array& a, + const std::vector& axes, + StreamOrDevice s /* = {} */) { + auto ndim = static_cast(a.ndim()); + Shape start(ndim, 0); + Shape stop = a.shape(); + Shape strides(ndim, 1); + for (auto ax : axes) { + int axis = normalize_axis_index(ax, ndim, "[flip] "); + start[axis] = a.shape(axis) - 1; + stop[axis] = -a.shape(axis) - 1; + strides[axis] = -1; + } + return slice(a, std::move(start), std::move(stop), std::move(strides), s); +} + +array flip(const array& a, int axis, StreamOrDevice s /* = {} */) { + return flip(a, std::vector{axis}, s); +} + +array flip(const array& a, StreamOrDevice s /* = {} */) { + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + return flip(a, axes, s); +} + // Slice helper namespace { @@ -1148,6 +1175,32 @@ split(const array& a, int num_splits, StreamOrDevice s /* = {} */) { return split(a, num_splits, 0, to_stream(s)); } +std::vector +unstack(const array& a, int axis, StreamOrDevice s /* = {} */) { + auto ndim = static_cast(a.ndim()); + auto ax = axis < 0 ? axis + ndim : axis; + if (ax < 0 || ax >= ndim) { + std::ostringstream msg; + msg << "[unstack] Invalid axis " << axis << " for array with " << ndim + << " dimensions."; + throw std::invalid_argument(msg.str()); + } + auto n = a.shape(ax); + std::vector res; + res.reserve(n); + if (n == 0) { + return res; + } + for (auto& part : split(a, n, ax, s)) { + res.push_back(squeeze(part, ax, s)); + } + return res; +} + +std::vector unstack(const array& a, StreamOrDevice s /* = {} */) { + return unstack(a, 0, s); +} + std::vector meshgrid( const std::vector& arrays, bool sparse /* = false */, diff --git a/mlx/ops.h b/mlx/ops.h index 208964d1aa..cd87aa4cd0 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -171,6 +171,16 @@ MLX_API array expand_dims( /** Add a singleton dimension at the given axis. */ MLX_API array expand_dims(const array& a, int axis, StreamOrDevice s = {}); +/** Reverse the order of the elements along the given axes. */ +MLX_API array +flip(const array& a, const std::vector& axes, StreamOrDevice s = {}); + +/** Reverse the order of the elements along the given axis. */ +MLX_API array flip(const array& a, int axis, StreamOrDevice s = {}); + +/** Reverse the order of the elements along all axes. */ +MLX_API array flip(const array& a, StreamOrDevice s = {}); + /** Slice an array. */ MLX_API array slice( const array& a, @@ -306,6 +316,11 @@ split(const array& a, const Shape& indices, int axis, StreamOrDevice s = {}); MLX_API std::vector split(const array& a, const Shape& indices, StreamOrDevice s = {}); +/** Split an array into a sequence of arrays along an axis, removing it. */ +MLX_API std::vector +unstack(const array& a, int axis, StreamOrDevice s = {}); +MLX_API std::vector unstack(const array& a, StreamOrDevice s = {}); + /** A vector of coordinate arrays from coordinate vectors. */ MLX_API std::vector meshgrid( const std::vector& arrays, diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 198a4861d9..5df47329a2 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -177,6 +177,60 @@ void init_ops(nb::module_& m) { Returns: array: The output array with size one axes removed. )pbdoc"); + m.def( + "flip", + [](const mx::array& a, const IntOrVec& v, const mx::StreamOrDevice& s) { + if (std::holds_alternative(v)) { + return mx::flip(a, s); + } else if (auto pv = std::get_if(&v); pv) { + return mx::flip(a, *pv, s); + } else { + return mx::flip(a, std::get>(v), s); + } + }, + nb::arg(), + "axis"_a = nb::none(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def flip(a: array, /, axis: Union[None, int, Sequence[int]] = None, " + "*, stream: Union[None, Stream, Device] = None) -> array"), + R"pbdoc( + Reverse the order of elements along the given axis. + + Args: + a (array): Input array. + axis (int or tuple(int), optional): Axis or axes to flip over. + Defaults to ``None`` in which case all axes are flipped. + + Returns: + array: The flipped array. + )pbdoc"); + m.def( + "unstack", + [](const mx::array& a, int axis, mx::StreamOrDevice s) { + return mx::unstack(a, axis, s); + }, + nb::arg(), + nb::kw_only(), + "axis"_a = 0, + "stream"_a = nb::none(), + nb::sig( + "def unstack(x: array, /, *, axis: int = 0, stream: Union[None, " + "Stream, Device] = None) -> list[array]"), + R"pbdoc( + Split an array into a sequence of arrays along the given axis. + + The inverse of :func:`stack`. The given axis is removed from each of + the returned arrays. + + Args: + x (array): Input array. + axis (int, optional): Axis along which to unstack. Default: ``0``. + + Returns: + list(array): A list of arrays, one for each index along ``axis``. + )pbdoc"); m.def( "expand_dims", [](const mx::array& a, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 484ae47c3a..c58b840e11 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1374,6 +1374,42 @@ def test_split(self): self.assertEqual(y.tolist(), [1, 2, 3, 4]) self.assertEqual(z.tolist(), [5, 6, 7]) + def test_flip(self): + a_np = np.arange(6).reshape(2, 3) + a = mx.array(a_np) + for axis in [None, 0, 1, -1, (0, 1)]: + self.assertTrue( + np.array_equal(mx.flip(a, axis=axis), np.flip(a_np, axis=axis)), + msg=f"axis={axis}", + ) + + # 1D + b_np = np.array([1, 2, 3, 4]) + self.assertTrue(np.array_equal(mx.flip(mx.array(b_np)), np.flip(b_np))) + + with self.assertRaises(ValueError): + mx.flip(a, axis=2) + + def test_unstack(self): + a_np = np.arange(6).reshape(3, 2) + a = mx.array(a_np) + for axis in [0, 1, -1]: + parts = mx.unstack(a, axis=axis) + expected = np.unstack(a_np, axis=axis) + self.assertEqual(len(parts), len(expected)) + for p, e in zip(parts, expected): + self.assertTrue(np.array_equal(p, e)) + self.assertEqual(p.shape, e.shape) + + # Default axis is 0. + self.assertTrue(np.array_equal(mx.unstack(a)[1], np.unstack(a_np)[1])) + + # stack is the inverse of unstack. + self.assertTrue(mx.array_equal(mx.stack(mx.unstack(a, axis=1), axis=1), a)) + + with self.assertRaises(ValueError): + mx.unstack(a, axis=2) + def test_split_invalid_num_splits(self): """Regression: split with num_splits <= 0 should raise, not crash.""" a = mx.arange(6) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index c2230f77cf..8b277c5580 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -590,6 +590,46 @@ TEST_CASE("test split") { CHECK(array_equal(out[3], array({2, 3, 4})).item()); } +TEST_CASE("test flip") { + array x = array({1, 2, 3, 4}); + CHECK(array_equal(flip(x), array({4, 3, 2, 1})).item()); + + x = array({0, 1, 2, 3, 4, 5}, {2, 3}); + CHECK( + array_equal(flip(x, 0), array({3, 4, 5, 0, 1, 2}, {2, 3})).item()); + CHECK( + array_equal(flip(x, 1), array({2, 1, 0, 5, 4, 3}, {2, 3})).item()); + CHECK( + array_equal(flip(x, -1), array({2, 1, 0, 5, 4, 3}, {2, 3})).item()); + // No axes -> flip all. + CHECK(array_equal(flip(x), array({5, 4, 3, 2, 1, 0}, {2, 3})).item()); + CHECK(array_equal( + flip(x, std::vector{0, 1}), array({5, 4, 3, 2, 1, 0}, {2, 3})) + .item()); + + CHECK_THROWS(flip(x, 2)); +} + +TEST_CASE("test unstack") { + array x = array({0, 1, 2, 3, 4, 5}, {3, 2}); + auto out = unstack(x); + CHECK_EQ(out.size(), 3); + CHECK(array_equal(out[0], array({0, 1})).item()); + CHECK(array_equal(out[1], array({2, 3})).item()); + CHECK(array_equal(out[2], array({4, 5})).item()); + CHECK_EQ(out[0].shape(), Shape{2}); + + out = unstack(x, 1); + CHECK_EQ(out.size(), 2); + CHECK(array_equal(out[0], array({0, 2, 4})).item()); + CHECK(array_equal(out[1], array({1, 3, 5})).item()); + + // stack is the inverse of unstack. + CHECK(array_equal(stack(unstack(x, 1), 1), x).item()); + + CHECK_THROWS(unstack(x, 2)); +} + TEST_CASE("test swap and move axes") { // Test swapaxes array a(0.0);