Use SmallVector for shapes and strides (#2454)

* Use SmallVector for shapes and strides

* Convert SmallVector to tuple

python/src/array.cpp

@@ -15,6 +15,7 @@
 #include "python/src/buffer.h"
 #include "python/src/convert.h"
 #include "python/src/indexing.h"
+#include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
 #include "mlx/mlx.h"
@@ -303,7 +304,7 @@ void init_array(nb::module_& m) {
           R"pbdoc(The number of bytes in the array.)pbdoc")
       .def_prop_ro(
           "shape",
-          [](const mx::array& a) { return nb::tuple(nb::cast(a.shape())); },
+          [](const mx::array& a) { return nb::cast(a.shape()); },
           R"pbdoc(
           The shape of the array as a Python tuple.
 

python/src/distributed.cpp

@@ -9,7 +9,7 @@
 
 #include "mlx/distributed/distributed.h"
 #include "mlx/distributed/ops.h"
-
+#include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
 namespace mx = mlx::core;

python/src/export.cpp

@@ -10,6 +10,7 @@
 #include "mlx/array.h"
 #include "mlx/export.h"
 #include "mlx/graph_utils.h"
+#include "python/src/small_vector.h"
 #include "python/src/trees.h"
 
 namespace mx = mlx::core;

python/src/fast.cpp

@@ -8,10 +8,10 @@
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
 
-#include "python/src/utils.h"
-
 #include "mlx/fast.h"
 #include "mlx/ops.h"
+#include "python/src/small_vector.h"
+#include "python/src/utils.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;

python/src/fft.cpp

@@ -8,6 +8,7 @@
 
 #include "mlx/fft.h"
 #include "mlx/ops.h"
+#include "python/src/small_vector.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;

python/src/linalg.cpp

@@ -9,6 +9,7 @@
 #include <nanobind/stl/vector.h>
 
 #include "mlx/linalg.h"
+#include "python/src/small_vector.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;

python/src/load.cpp

@@ -12,6 +12,7 @@
 #include "mlx/ops.h"
 #include "mlx/utils.h"
 #include "python/src/load.h"
+#include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
 namespace mx = mlx::core;

python/src/metal.cpp

@@ -7,8 +7,10 @@
 #include <nanobind/stl/unordered_map.h>
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
+
 #include "mlx/backend/metal/metal.h"
 #include "mlx/memory.h"
+#include "python/src/small_vector.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;

python/src/ops.cpp

@@ -16,6 +16,7 @@
 #include "mlx/ops.h"
 #include "mlx/utils.h"
 #include "python/src/load.h"
+#include "python/src/small_vector.h"
 #include "python/src/utils.h"
 
 namespace mx = mlx::core;

python/src/random.cpp

@@ -7,10 +7,10 @@
 
 #include <chrono>
 
-#include "python/src/utils.h"
-
 #include "mlx/ops.h"
 #include "mlx/random.h"
+#include "python/src/small_vector.h"
+#include "python/src/utils.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;

python/src/small_vector.h

@@ -0,0 +1,77 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/small_vector.h"
+
+#include <nanobind/stl/detail/nb_list.h>
+
+NAMESPACE_BEGIN(NB_NAMESPACE)
+NAMESPACE_BEGIN(detail)
+
+template <typename Type, size_t Size, typename Alloc>
+struct type_caster<mlx::core::SmallVector<Type, Size, Alloc>> {
+  using List = mlx::core::SmallVector<Type, Size, Alloc>;
+  using Caster = make_caster<Type>;
+
+  NB_TYPE_CASTER(
+      List,
+      const_name(NB_TYPING_TUPLE "[") + make_caster<Type>::Name +
+          const_name(", ...]"))
+
+  bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
+    size_t size;
+    PyObject* temp;
+
+    // Will initialize 'size' and 'temp'. All return values and
+    // return parameters are zero/NULL in the case of a failure.
+    PyObject** o = seq_get(src.ptr(), &size, &temp);
+
+    value.clear();
+    value.reserve(size);
+
+    Caster caster;
+    bool success = o != nullptr;
+
+    flags = flags_for_local_caster<Type>(flags);
+
+    for (size_t i = 0; i < size; ++i) {
+      if (!caster.from_python(o[i], flags, cleanup) ||
+          !caster.template can_cast<Type>()) {
+        success = false;
+        break;
+      }
+
+      value.push_back(caster.operator cast_t<Type>());
+    }
+
+    Py_XDECREF(temp);
+
+    return success;
+  }
+
+  template <typename T>
+  static handle from_cpp(T&& src, rv_policy policy, cleanup_list* cleanup) {
+    object ret = steal(PyTuple_New(src.size()));
+
+    if (ret.is_valid()) {
+      Py_ssize_t index = 0;
+
+      for (auto&& value : src) {
+        handle h = Caster::from_cpp(forward_like_<T>(value), policy, cleanup);
+
+        if (!h.is_valid()) {
+          ret.reset();
+          break;
+        }
+
+        NB_TUPLE_SET_ITEM(ret.ptr(), index++, h.ptr());
+      }
+    }
+
+    return ret.release();
+  }
+};
+
+NAMESPACE_END(detail)
+NAMESPACE_END(NB_NAMESPACE)

python/src/transforms.cpp

@@ -20,6 +20,7 @@
 #include "mlx/transforms_impl.h"
 #include "mlx/utils.h"
 #include "python/src/mlx_func.h"
+#include "python/src/small_vector.h"
 #include "python/src/trees.h"
 
 namespace mx = mlx::core;