Add flip and unstack

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

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<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  auto ndim = static_cast<int>(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<int>{axis}, s);
+}
+
+array flip(const array& a, StreamOrDevice s /* = {} */) {
+  std::vector<int> 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<array>
+unstack(const array& a, int axis, StreamOrDevice s /* = {} */) {
+  auto ndim = static_cast<int>(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<array> 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<array> unstack(const array& a, StreamOrDevice s /* = {} */) {
+  return unstack(a, 0, s);
+}
+
 std::vector<array> meshgrid(
     const std::vector<array>& arrays,
     bool sparse /* = false */,

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<int>& 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<array>
 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<array>
+unstack(const array& a, int axis, StreamOrDevice s = {});
+MLX_API std::vector<array> unstack(const array& a, StreamOrDevice s = {});
+
 /** A vector of coordinate arrays from coordinate vectors. */
 MLX_API std::vector<array> meshgrid(
     const std::vector<array>& arrays,

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<std::monostate>(v)) {
+          return mx::flip(a, s);
+        } else if (auto pv = std::get_if<int>(&v); pv) {
+          return mx::flip(a, *pv, s);
+        } else {
+          return mx::flip(a, std::get<std::vector<int>>(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,

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)

tests/ops_tests.cpp

@@ -590,6 +590,46 @@ TEST_CASE("test split") {
   CHECK(array_equal(out[3], array({2, 3, 4})).item<bool>());
 }
 
+TEST_CASE("test flip") {
+  array x = array({1, 2, 3, 4});
+  CHECK(array_equal(flip(x), array({4, 3, 2, 1})).item<bool>());
+
+  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<bool>());
+  CHECK(
+      array_equal(flip(x, 1), array({2, 1, 0, 5, 4, 3}, {2, 3})).item<bool>());
+  CHECK(
+      array_equal(flip(x, -1), array({2, 1, 0, 5, 4, 3}, {2, 3})).item<bool>());
+  // No axes -> flip all.
+  CHECK(array_equal(flip(x), array({5, 4, 3, 2, 1, 0}, {2, 3})).item<bool>());
+  CHECK(array_equal(
+            flip(x, std::vector<int>{0, 1}), array({5, 4, 3, 2, 1, 0}, {2, 3}))
+            .item<bool>());
+
+  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<bool>());
+  CHECK(array_equal(out[1], array({2, 3})).item<bool>());
+  CHECK(array_equal(out[2], array({4, 5})).item<bool>());
+  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<bool>());
+  CHECK(array_equal(out[1], array({1, 3, 5})).item<bool>());
+
+  // stack is the inverse of unstack.
+  CHECK(array_equal(stack(unstack(x, 1), 1), x).item<bool>());
+
+  CHECK_THROWS(unstack(x, 2));
+}
+
 TEST_CASE("test swap and move axes") {
   // Test swapaxes
   array a(0.0);