Adding minmax method to MLX to follow Numpy/Jax API

docs/src/python/ops.rst

@@ -134,6 +134,7 @@ Operations
    meshgrid
    min
    minimum
+   minmax
    moveaxis
    multiply
    nan_to_num

mlx/ops.cpp

@@ -2657,6 +2657,31 @@ array min(
   return min(a, std::vector<int>{axis}, keepdims, s);
 }
 
+std::tuple<array, array>
+minmax(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return {min(a, axes, keepdims, s), max(a, axes, keepdims, s)};
+}
+
+std::tuple<array, array> minmax(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  return {min(a, axes, keepdims, s), max(a, axes, keepdims, s)};
+}
+
+std::tuple<array, array> minmax(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return {
+      min(a, std::vector<int>{axis}, keepdims, s),
+      max(a, std::vector<int>{axis}, keepdims, s)};
+}
+
 array argmin(const array& a, bool keepdims, StreamOrDevice s /* = {} */) {
   auto result = argmin(flatten(a, s), 0, true, s);
   if (keepdims) {

mlx/ops.h

@@ -786,6 +786,23 @@ min(const array& a,
 MLX_API array
 min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
 
+MLX_API std::tuple<array, array>
+minmax(const array& a, bool keepdims, StreamOrDevice s = {});
+inline std::tuple<array, array> minmax(const array& a, StreamOrDevice s = {}) {
+  return minmax(a, false, to_stream(s));
+}
+
+/** The minimum and maximum of the elements of an array along the given axes. */
+MLX_API std::tuple<array, array> minmax(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** The minimum and maximum of the elements of an array along the given axis. */
+MLX_API std::tuple<array, array>
+minmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
 /** Returns the Hanning window of size M. */
 MLX_API array hanning(int M, StreamOrDevice s = {});
 

python/src/ops.cpp

@@ -2738,6 +2738,37 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The output array with the corresponding axes reduced.
       )pbdoc");
+  m.def(
+      "minmax",
+      [](const mx::array& a,
+         const IntOrVec& axis,
+         bool keepdims,
+         mx::StreamOrDevice s) -> std::tuple<mx::array, mx::array> {
+        auto axes = get_reduce_axes(axis, a.ndim());
+        return mx::minmax(a, axes, keepdims, s);
+      },
+      nb::arg(),
+      "axis"_a = nb::none(),
+      "keepdims"_a = false,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def minmax(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> tuple[array, array]"),
+      R"pbdoc(
+        Return the minimum and maximum over the given axes.
+
+        Equivalent to calling :func:`min` and :func:`max` separately.
+
+        Args:
+            a (array): Input array.
+            axis (int or list(int), optional): Axis or axes to reduce over.
+              If unspecified, reduce over the entire array.
+            keepdims (bool, optional): Keep reduced axes as singleton
+              dimensions. Default: ``False``.
+
+        Returns:
+            tuple(array, array): The minimum and maximum arrays, in that order.
+      )pbdoc");
   m.def(
       "logcumsumexp",
       [](const mx::array& a,

python/tests/test_reduce.py

@@ -367,6 +367,21 @@ def test_init_reduce_large_output(self):
         mx.eval(out)
         self.assertTrue(mx.all(out).item())
 
+    def test_minmax(self):
+        x = mx.array([[3.0, 1.0, 4.0], [2.0, 6.0, 5.0]])
+        out = mx.minmax(x)
+        self.assertIsInstance(out, tuple)
+        self.assertEqual(len(out), 2)
+        self.assertEqual(tuple(a.item() for a in out), (1.0, 6.0))
+        for axis in (None, 0, (0, 1)):
+            for keepdims in (False, True):
+                with self.subTest(axis=axis, keepdims=keepdims):
+                    out = mx.minmax(x, axis=axis, keepdims=keepdims)
+                    for actual, op in zip(out, (mx.min, mx.max)):
+                        expected = op(x, axis=axis, keepdims=keepdims)
+                        self.assertEqual(actual.shape, expected.shape)
+                        self.assertTrue(mx.array_equal(actual, expected))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner(failfast=True)