Bump (#869)

* bump

* fix none in a few ops

.circleci/config.yml

@@ -31,8 +31,7 @@ jobs:
           name: Install dependencies
           command: |
             pip install --upgrade cmake
-            pip install --upgrade pybind11[global]
-            pip install pybind11-stubgen
+            pip install git+https://github.com/wjakob/nanobind.git@4148debcf91f5ccab0c3b8d67b5c3cabd61f407f
             pip install numpy
             sudo apt-get update
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
@@ -44,7 +43,8 @@ jobs:
       - run:
           name: Generate package stubs
           command: |
-            python3 setup.py generate_stubs
+            echo "stubs"
+            python -m nanobind.stubgen -m mlx.core -r -O python
       - run:
           name: Run Python tests
           command: |
@@ -80,8 +80,7 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install --upgrade pybind11[global]
-            pip install pybind11-stubgen
+            pip install git+https://github.com/wjakob/nanobind.git@4148debcf91f5ccab0c3b8d67b5c3cabd61f407f
             pip install numpy
             pip install torch
             pip install tensorflow
@@ -95,7 +94,7 @@ jobs:
           name: Generate package stubs
           command: |
             source env/bin/activate
-            python setup.py generate_stubs
+            python -m nanobind.stubgen -m mlx.core -r -O python
       - run:
           name: Run Python tests
           command: |
@@ -144,9 +143,8 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install --upgrade pybind11[global]
+            pip install git+https://github.com/wjakob/nanobind.git@4148debcf91f5ccab0c3b8d67b5c3cabd61f407f
             pip install --upgrade setuptools
-            pip install pybind11-stubgen
             pip install numpy
             pip install twine
             pip install build
@@ -161,7 +159,7 @@ jobs:
           name: Generate package stubs
           command: |
             source env/bin/activate
-            python setup.py generate_stubs
+            python -m nanobind.stubgen -m mlx.core -r -O python
       - run:
           name: Build Python package
           command: |
@@ -209,9 +207,8 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install --upgrade pybind11[global]
+            pip install git+https://github.com/wjakob/nanobind.git@4148debcf91f5ccab0c3b8d67b5c3cabd61f407f
             pip install --upgrade setuptools
-            pip install pybind11-stubgen
             pip install numpy
             pip install auditwheel
             pip install patchelf
@@ -219,7 +216,7 @@ jobs:
             << parameters.extra_env >> \
               CMAKE_BUILD_PARALLEL_LEVEL="" \
               pip install . -v
-            python setup.py generate_stubs
+            python -m nanobind.stubgen -m mlx.core -r -O python
             << parameters.extra_env >> \
               CMAKE_BUILD_PARALLEL_LEVEL="" \
               python -m build --wheel

CMakeLists.txt

@@ -18,7 +18,7 @@ option(MLX_BUILD_METAL "Build metal backend" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.7.0)
+  set(MLX_VERSION 0.8.0)
 endif()
 
 # --------------------- Processor tests -------------------------
@@ -146,8 +146,12 @@ target_include_directories(
 
 if (MLX_BUILD_PYTHON_BINDINGS)
   message(STATUS "Building Python bindings.")
-  find_package(Python COMPONENTS Interpreter Development)
-  find_package(pybind11 CONFIG REQUIRED)
+  find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE NB_DIR)
+  list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
+  find_package(nanobind CONFIG REQUIRED)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 

docs/src/conf.py

@@ -29,8 +29,8 @@ autosummary_generate = True
 autosummary_filename_map = {"mlx.core.Stream": "stream_class"}
 
 intersphinx_mapping = {
-    "https://docs.python.org/3": None,
-    "https://numpy.org/doc/stable/": None,
+    "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable/", None),
 }
 
 templates_path = ["_templates"]
@@ -59,3 +59,14 @@ html_theme_options = {
 # -- Options for HTMLHelp output ---------------------------------------------
 
 htmlhelp_basename = "mlx_doc"
+
+
+def setup(app):
+    wrapped = app.registry.documenters["function"].can_document_member
+
+    def nanobind_function_patch(member: Any, *args, **kwargs) -> bool:
+        return "nanobind.nb_func" in str(type(member)) or wrapped(
+            member, *args, **kwargs
+        )
+
+    app.registry.documenters["function"].can_document_member = nanobind_function_patch

docs/src/install.rst

@@ -70,16 +70,13 @@ To build and install the MLX python library from source, first, clone MLX from
 
    git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
 
-Make sure that you have `pybind11 <https://pybind11.readthedocs.io/en/stable/index.html>`_
-installed. You can install ``pybind11`` with ``pip``, ``brew`` or ``conda`` as follows:
+Install `nanobind <https://nanobind.readthedocs.io/en/latest/>`_ with:
 
 .. code-block:: shell
 
-    pip install "pybind11[global]"
-    conda install pybind11
-    brew install pybind11
+    pip install git+https://github.com/wjakob/nanobind.git
 
-Then simply build and install it using pip:
+Then simply build and install MLX using pip:
 
 .. code-block:: shell
 

mlx/array.cpp

@@ -36,22 +36,11 @@ array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
   init(&cval);
 }
 
-array::array(
-    const std::vector<int>& shape,
-    Dtype dtype,
-    std::shared_ptr<Primitive> primitive,
-    const std::vector<array>& inputs)
-    : array_desc_(std::make_shared<ArrayDesc>(
-          shape,
-          dtype,
-          std::move(primitive),
-          inputs)) {}
-
 array::array(
     std::vector<int> shape,
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
-    std::vector<array>&& inputs)
+    std::vector<array> inputs)
     : array_desc_(std::make_shared<ArrayDesc>(
           std::move(shape),
           dtype,
@@ -92,10 +81,10 @@ array::array(std::initializer_list<int> data, Dtype dtype)
 /* Build an array from a shared buffer */
 array::array(
     allocator::Buffer data,
-    const std::vector<int>& shape,
+    std::vector<int> shape,
     Dtype dtype,
     deleter_t deleter)
-    : array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   set_data(data, deleter);
 }
 
@@ -181,31 +170,16 @@ void array::move_shared_buffer(array other) {
   move_shared_buffer(other, other.strides(), other.flags(), other.data_size());
 }
 
-array::ArrayDesc::ArrayDesc(const std::vector<int>& shape, Dtype dtype)
-    : shape(shape), dtype(dtype) {
-  std::tie(size, strides) = cum_prod(shape);
-}
-
-array::ArrayDesc::ArrayDesc(
-    const std::vector<int>& shape,
-    Dtype dtype,
-    std::shared_ptr<Primitive> primitive,
-    const std::vector<array>& inputs)
-    : shape(shape),
-      dtype(dtype),
-      primitive(std::move(primitive)),
-      inputs(inputs) {
+array::ArrayDesc::ArrayDesc(std::vector<int> shape, Dtype dtype)
+    : shape(std::move(shape)), dtype(dtype) {
   std::tie(size, strides) = cum_prod(this->shape);
-  for (auto& in : this->inputs) {
-    is_tracer |= in.is_tracer();
-  }
 }
 
 array::ArrayDesc::ArrayDesc(
-    std::vector<int>&& shape,
+    std::vector<int> shape,
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
-    std::vector<array>&& inputs)
+    std::vector<array> inputs)
     : shape(std::move(shape)),
       dtype(dtype),
       primitive(std::move(primitive)),

mlx/array.h

@@ -1,5 +1,6 @@
 // Copyright © 2023 Apple Inc.
 #pragma once
+
 #include <algorithm>
 #include <cstdint>
 #include <functional>
@@ -31,7 +32,7 @@ class array {
   template <typename It>
   array(
       It data,
-      const std::vector<int>& shape,
+      std::vector<int> shape,
       Dtype dtype =
           TypeToDtype<typename std::iterator_traits<It>::value_type>());
 
@@ -47,13 +48,13 @@ class array {
   template <typename T>
   array(
       std::initializer_list<T> data,
-      const std::vector<int>& shape,
+      std::vector<int> shape,
       Dtype dtype = TypeToDtype<T>());
 
   /* Build an array from a buffer */
   array(
       allocator::Buffer data,
-      const std::vector<int>& shape,
+      std::vector<int> shape,
       Dtype dtype,
       deleter_t deleter = allocator::free);
 
@@ -172,17 +173,11 @@ class array {
    * API may change.
    */
 
-  array(
-      const std::vector<int>& shape,
-      Dtype dtype,
-      std::shared_ptr<Primitive> primitive,
-      const std::vector<array>& inputs);
-
   array(
       std::vector<int> shape,
       Dtype dtype,
       std::shared_ptr<Primitive> primitive,
-      std::vector<array>&& inputs);
+      std::vector<array> inputs);
 
   static std::vector<array> make_arrays(
       const std::vector<std::vector<int>>& shapes,
@@ -362,7 +357,7 @@ class array {
     std::vector<size_t> strides;
     size_t size;
     Dtype dtype;
-    std::shared_ptr<Primitive> primitive{nullptr};
+    std::shared_ptr<Primitive> primitive;
 
     // Indicates an array is being used in a graph transform
     // and should not be detached from the graph
@@ -370,7 +365,7 @@ class array {
 
     // This is a shared pointer so that *different* arrays
     // can share the underlying data buffer.
-    std::shared_ptr<Data> data{nullptr};
+    std::shared_ptr<Data> data;
 
     // Properly offset data pointer
     void* data_ptr{nullptr};
@@ -390,26 +385,20 @@ class array {
     // The arrays position in the output list
     uint32_t position{0};
 
-    explicit ArrayDesc(const std::vector<int>& shape, Dtype dtype);
+    explicit ArrayDesc(std::vector<int> shape, Dtype dtype);
 
     explicit ArrayDesc(
-        const std::vector<int>& shape,
+        std::vector<int> shape,
         Dtype dtype,
         std::shared_ptr<Primitive> primitive,
-        const std::vector<array>& inputs);
-
-    explicit ArrayDesc(
-        std::vector<int>&& shape,
-        Dtype dtype,
-        std::shared_ptr<Primitive> primitive,
-        std::vector<array>&& inputs);
+        std::vector<array> inputs);
   };
 
   // The ArrayDesc contains the details of the materialized array including the
   // shape, strides, the data type. It also includes
   // the primitive which knows how to compute the array's data from its inputs
   // and the list of array's inputs for the primitive.
-  std::shared_ptr<ArrayDesc> array_desc_{nullptr};
+  std::shared_ptr<ArrayDesc> array_desc_;
 };
 
 template <typename T>
@@ -421,9 +410,9 @@ array::array(T val, Dtype dtype /* = TypeToDtype<T>() */)
 template <typename It>
 array::array(
   It data,
-  const std::vector<int>& shape,
+  std::vector<int> shape,
   Dtype dtype /* = TypeToDtype<typename std::iterator_traits<It>::value_type>() */) :
-    array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
+    array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   init(data);
 }
 
@@ -440,9 +429,9 @@ array::array(
 template <typename T>
 array::array(
     std::initializer_list<T> data,
-    const std::vector<int>& shape,
+    std::vector<int> shape,
     Dtype dtype /* = TypeToDtype<T>() */)
-    : array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   if (data.size() != size()) {
     throw std::invalid_argument(
         "Data size and provided shape mismatch in array construction.");

mlx/backend/accelerate/primitives.cpp

@@ -69,11 +69,13 @@ DEFAULT(Select)
 DEFAULT(Sigmoid)
 DEFAULT(Sign)
 DEFAULT(Slice)
+DEFAULT(SliceUpdate)
 DEFAULT_MULTI(Split)
 DEFAULT(Sort)
 DEFAULT(StopGradient)
 DEFAULT_MULTI(SVD)
 DEFAULT(Transpose)
+DEFAULT(Inverse)
 
 void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);

mlx/backend/common/CMakeLists.txt

@@ -44,7 +44,6 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
@@ -54,6 +53,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/qrf.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/svd.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/inverse.cpp
   ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
 )
 

mlx/backend/common/copy.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <numeric>
 
@@ -25,121 +25,196 @@ void copy_vector(const array& src, array& dst) {
   std::copy(src_ptr, src_ptr + src.data_size(), dst_ptr);
 }
 
-template <typename SrcT, typename DstT>
-void copy_general_dim1(const array& src, array& dst) {
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general_dim1(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    int64_t i_offset) {
   const SrcT* src_ptr = src.data<SrcT>();
   DstT* dst_ptr = dst.data<DstT>();
-  size_t src_idx = 0;
-  size_t dst_idx = 0;
-  for (size_t i = 0; i < src.shape()[0]; ++i) {
+  stride_t src_idx = i_offset;
+  stride_t dst_idx = 0;
+  for (int i = 0; i < data_shape[0]; ++i) {
     dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-    src_idx += src.strides()[0];
+    src_idx += i_strides[0];
   }
 }
 
 template <typename SrcT, typename DstT>
-void copy_general_dim2(const array& src, array& dst) {
+inline void copy_general_dim1(const array& src, array& dst) {
+  return copy_general_dim1<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), 0);
+}
+
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general_dim2(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    int64_t i_offset) {
   const SrcT* src_ptr = src.data<SrcT>();
   DstT* dst_ptr = dst.data<DstT>();
-  size_t src_idx = 0;
-  size_t dst_idx = 0;
-  for (size_t i = 0; i < src.shape()[0]; ++i) {
-    for (size_t j = 0; j < src.shape()[1]; ++j) {
+  stride_t src_idx = i_offset;
+  stride_t dst_idx = 0;
+  for (int i = 0; i < data_shape[0]; ++i) {
+    for (int j = 0; j < data_shape[1]; ++j) {
       dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-      src_idx += src.strides()[1];
+      src_idx += i_strides[1];
     }
-    src_idx += src.strides()[0] - src.strides()[1] * src.shape()[1];
+    src_idx += i_strides[0] - i_strides[1] * data_shape[1];
   }
 }
 
 template <typename SrcT, typename DstT>
-void copy_general_dim3(const array& src, array& dst) {
+inline void copy_general_dim2(const array& src, array& dst) {
+  return copy_general_dim2<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), 0);
+}
+
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general_dim3(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    int64_t i_offset) {
   const SrcT* src_ptr = src.data<SrcT>();
   DstT* dst_ptr = dst.data<DstT>();
-  size_t src_idx = 0;
-  size_t dst_idx = 0;
-  for (size_t i = 0; i < src.shape()[0]; ++i) {
-    for (size_t j = 0; j < src.shape()[1]; ++j) {
-      for (size_t k = 0; k < src.shape()[2]; ++k) {
+  stride_t src_idx = i_offset;
+  stride_t dst_idx = 0;
+  for (int i = 0; i < data_shape[0]; ++i) {
+    for (int j = 0; j < data_shape[1]; ++j) {
+      for (int k = 0; k < data_shape[2]; ++k) {
         dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-        src_idx += src.strides()[2];
+        src_idx += i_strides[2];
       }
-      src_idx += src.strides()[1] - src.strides()[2] * src.shape()[2];
+      src_idx += i_strides[1] - i_strides[2] * data_shape[2];
     }
-    src_idx += src.strides()[0] - src.strides()[1] * src.shape()[1];
+    src_idx += i_strides[0] - i_strides[1] * data_shape[1];
   }
 }
 
 template <typename SrcT, typename DstT>
-void copy_general_dim4(const array& src, array& dst) {
+inline void copy_general_dim3(const array& src, array& dst) {
+  return copy_general_dim3<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), 0);
+}
+
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general_dim4(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    int64_t i_offset) {
   const SrcT* src_ptr = src.data<SrcT>();
   DstT* dst_ptr = dst.data<DstT>();
-  size_t src_idx = 0;
-  size_t dst_idx = 0;
-  for (size_t i = 0; i < src.shape()[0]; ++i) {
-    for (size_t j = 0; j < src.shape()[1]; ++j) {
-      for (size_t k = 0; k < src.shape()[2]; ++k) {
-        for (size_t ii = 0; ii < src.shape()[3]; ++ii) {
+  stride_t src_idx = i_offset;
+  stride_t dst_idx = 0;
+  for (int i = 0; i < data_shape[0]; ++i) {
+    for (int j = 0; j < data_shape[1]; ++j) {
+      for (int k = 0; k < data_shape[2]; ++k) {
+        for (int ii = 0; ii < data_shape[3]; ++ii) {
           dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-          src_idx += src.strides()[3];
+          src_idx += i_strides[3];
         }
-        src_idx += src.strides()[2] - src.strides()[3] * src.shape()[3];
+        src_idx += i_strides[2] - i_strides[3] * data_shape[3];
       }
-      src_idx += src.strides()[1] - src.strides()[2] * src.shape()[2];
+      src_idx += i_strides[1] - i_strides[2] * data_shape[2];
     }
-    src_idx += src.strides()[0] - src.strides()[1] * src.shape()[1];
+    src_idx += i_strides[0] - i_strides[1] * data_shape[1];
   }
 }
 
 template <typename SrcT, typename DstT>
-void copy_general(const array& src, array& dst) {
+inline void copy_general_dim4(const array& src, array& dst) {
+  return copy_general_dim4<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), 0);
+}
+
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    int64_t i_offset) {
   switch (src.ndim()) {
     case 1:
-      copy_general_dim1<SrcT, DstT>(src, dst);
+      copy_general_dim1<SrcT, DstT, stride_t>(
+          src, dst, data_shape, i_strides, i_offset);
       return;
     case 2:
-      copy_general_dim2<SrcT, DstT>(src, dst);
+      copy_general_dim2<SrcT, DstT, stride_t>(
+          src, dst, data_shape, i_strides, i_offset);
       return;
     case 3:
-      copy_general_dim3<SrcT, DstT>(src, dst);
+      copy_general_dim3<SrcT, DstT, stride_t>(
+          src, dst, data_shape, i_strides, i_offset);
       return;
     case 4:
-      copy_general_dim4<SrcT, DstT>(src, dst);
+      copy_general_dim4<SrcT, DstT, stride_t>(
+          src, dst, data_shape, i_strides, i_offset);
       return;
   }
 
-  auto src_ptr = src.data<SrcT>();
+  auto src_ptr = src.data<SrcT>() + i_offset;
   auto dst_ptr = dst.data<DstT>();
   for (size_t i = 0; i < dst.size(); ++i) {
-    size_t src_elem = elem_to_loc(i, src.shape(), src.strides());
+    stride_t src_elem = elem_to_loc(i, data_shape, i_strides);
     dst_ptr[i] = static_cast<DstT>(src_ptr[src_elem]);
   }
 }
 
-template <typename SrcT, typename DstT, int D>
+template <typename SrcT, typename DstT>
+inline void copy_general(const array& src, array& dst) {
+  return copy_general<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), 0);
+}
+
+template <typename SrcT, typename DstT, typename stride_t>
+inline void copy_general(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset) {
+  return copy_general<SrcT, DstT, stride_t>(
+      src, dst, data_shape, i_strides, i_offset);
+}
+
+template <typename SrcT, typename DstT, typename stride_t, int D>
 inline void copy_general_general_dims(
     const array& src,
     array& dst,
-    size_t offset_src,
-    size_t offset_dst) {
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    stride_t i_offset,
+    stride_t o_offset) {
   if constexpr (D > 1) {
     int axis = src.ndim() - D;
-    auto stride_src = src.strides()[axis];
-    auto stride_dst = dst.strides()[axis];
-    auto N = src.shape(axis);
+    auto stride_src = i_strides[axis];
+    auto stride_dst = o_strides[axis];
+    auto N = data_shape[axis];
     for (int i = 0; i < N; i++) {
-      copy_general_general_dims<SrcT, DstT, D - 1>(
-          src, dst, offset_src, offset_dst);
-      offset_src += stride_src;
-      offset_dst += stride_dst;
+      copy_general_general_dims<SrcT, DstT, stride_t, D - 1>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+      i_offset += stride_src;
+      o_offset += stride_dst;
     }
   } else {
     int axis = src.ndim() - 1;
-    auto stride_src = src.strides()[axis];
-    auto stride_dst = dst.strides()[axis];
-    auto N = src.shape(axis);
-    const SrcT* src_ptr = src.data<SrcT>() + offset_src;
-    DstT* dst_ptr = dst.data<DstT>() + offset_dst;
+    auto stride_src = i_strides[axis];
+    auto stride_dst = o_strides[axis];
+    auto N = data_shape[axis];
+    const SrcT* src_ptr = src.data<SrcT>() + i_offset;
+    DstT* dst_ptr = dst.data<DstT>() + o_offset;
     for (int i = 0; i < N; i++) {
       *dst_ptr = static_cast<DstT>(*src_ptr);
       src_ptr += stride_src;
@@ -148,37 +223,56 @@ inline void copy_general_general_dims(
   }
 }
 
-template <typename SrcT, typename DstT>
-void copy_general_general(const array& src, array& dst) {
+template <typename SrcT, typename DstT, typename stride_t>
+void copy_general_general(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    stride_t i_offset,
+    stride_t o_offset) {
   switch (src.ndim()) {
     case 1:
-      copy_general_general_dims<SrcT, DstT, 1>(src, dst, 0, 0);
+      copy_general_general_dims<SrcT, DstT, stride_t, 1>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
       return;
     case 2:
-      copy_general_general_dims<SrcT, DstT, 2>(src, dst, 0, 0);
+      copy_general_general_dims<SrcT, DstT, stride_t, 2>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
       return;
     case 3:
-      copy_general_general_dims<SrcT, DstT, 3>(src, dst, 0, 0);
+      copy_general_general_dims<SrcT, DstT, stride_t, 3>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
       return;
     case 4:
-      copy_general_general_dims<SrcT, DstT, 4>(src, dst, 0, 0);
+      copy_general_general_dims<SrcT, DstT, stride_t, 4>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
       return;
     case 5:
-      copy_general_general_dims<SrcT, DstT, 5>(src, dst, 0, 0);
+      copy_general_general_dims<SrcT, DstT, stride_t, 5>(
+          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
       return;
   }
 
   int size = std::accumulate(
-      src.shape().begin() - 5, src.shape().end(), 1, std::multiplies<int>());
+      data_shape.begin() - 5, data_shape.end(), 1, std::multiplies<int>());
   for (int i = 0; i < src.size(); i += size) {
-    size_t offset_src = elem_to_loc(i, src.shape(), src.strides());
-    size_t offset_dst = elem_to_loc(i, dst.shape(), dst.strides());
-    copy_general_general_dims<SrcT, DstT, 5>(src, dst, offset_src, offset_dst);
+    stride_t src_offset = i_offset + elem_to_loc(i, data_shape, i_strides);
+    stride_t dst_offset = o_offset + elem_to_loc(i, dst.shape(), o_strides);
+    copy_general_general_dims<SrcT, DstT, stride_t, 5>(
+        src, dst, data_shape, i_strides, o_strides, src_offset, dst_offset);
   }
 }
 
 template <typename SrcT, typename DstT>
-void copy(const array& src, array& dst, CopyType ctype) {
+inline void copy_general_general(const array& src, array& dst) {
+  return copy_general_general<SrcT, DstT, size_t>(
+      src, dst, src.shape(), src.strides(), dst.strides(), 0, 0);
+}
+
+template <typename SrcT, typename DstT, typename... Args>
+void copy(const array& src, array& dst, CopyType ctype, Args... args) {
   switch (ctype) {
     case CopyType::Scalar:
       copy_single<SrcT, DstT>(src, dst);
@@ -187,54 +281,103 @@ void copy(const array& src, array& dst, CopyType ctype) {
       copy_vector<SrcT, DstT>(src, dst);
       return;
     case CopyType::General:
-      copy_general<SrcT, DstT>(src, dst);
+      copy_general<SrcT, DstT>(src, dst, args...);
       return;
     case CopyType::GeneralGeneral:
-      copy_general_general<SrcT, DstT>(src, dst);
+      copy_general_general<SrcT, DstT>(src, dst, args...);
   }
 }
 
-template <typename SrcT>
-void copy(const array& src, array& dst, CopyType ctype) {
+template <typename SrcT, typename... Args>
+void copy(const array& src, array& dst, CopyType ctype, Args... args) {
   switch (dst.dtype()) {
     case bool_:
-      copy<SrcT, bool>(src, dst, ctype);
+      copy<SrcT, bool>(src, dst, ctype, args...);
       break;
     case uint8:
-      copy<SrcT, uint8_t>(src, dst, ctype);
+      copy<SrcT, uint8_t>(src, dst, ctype, args...);
       break;
     case uint16:
-      copy<SrcT, uint16_t>(src, dst, ctype);
+      copy<SrcT, uint16_t>(src, dst, ctype, args...);
       break;
     case uint32:
-      copy<SrcT, uint32_t>(src, dst, ctype);
+      copy<SrcT, uint32_t>(src, dst, ctype, args...);
       break;
     case uint64:
-      copy<SrcT, uint64_t>(src, dst, ctype);
+      copy<SrcT, uint64_t>(src, dst, ctype, args...);
       break;
     case int8:
-      copy<SrcT, int8_t>(src, dst, ctype);
+      copy<SrcT, int8_t>(src, dst, ctype, args...);
       break;
     case int16:
-      copy<SrcT, int16_t>(src, dst, ctype);
+      copy<SrcT, int16_t>(src, dst, ctype, args...);
       break;
     case int32:
-      copy<SrcT, int32_t>(src, dst, ctype);
+      copy<SrcT, int32_t>(src, dst, ctype, args...);
       break;
     case int64:
-      copy<SrcT, int64_t>(src, dst, ctype);
+      copy<SrcT, int64_t>(src, dst, ctype, args...);
       break;
     case float16:
-      copy<SrcT, float16_t>(src, dst, ctype);
+      copy<SrcT, float16_t>(src, dst, ctype, args...);
       break;
     case float32:
-      copy<SrcT, float>(src, dst, ctype);
+      copy<SrcT, float>(src, dst, ctype, args...);
       break;
     case bfloat16:
-      copy<SrcT, bfloat16_t>(src, dst, ctype);
+      copy<SrcT, bfloat16_t>(src, dst, ctype, args...);
       break;
     case complex64:
-      copy<SrcT, complex64_t>(src, dst, ctype);
+      copy<SrcT, complex64_t>(src, dst, ctype, args...);
+      break;
+  }
+}
+
+template <typename... Args>
+inline void copy_inplace_dispatch(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Args... args) {
+  switch (src.dtype()) {
+    case bool_:
+      copy<bool>(src, dst, ctype, args...);
+      break;
+    case uint8:
+      copy<uint8_t>(src, dst, ctype, args...);
+      break;
+    case uint16:
+      copy<uint16_t>(src, dst, ctype, args...);
+      break;
+    case uint32:
+      copy<uint32_t>(src, dst, ctype, args...);
+      break;
+    case uint64:
+      copy<uint64_t>(src, dst, ctype, args...);
+      break;
+    case int8:
+      copy<int8_t>(src, dst, ctype, args...);
+      break;
+    case int16:
+      copy<int16_t>(src, dst, ctype, args...);
+      break;
+    case int32:
+      copy<int32_t>(src, dst, ctype, args...);
+      break;
+    case int64:
+      copy<int64_t>(src, dst, ctype, args...);
+      break;
+    case float16:
+      copy<float16_t>(src, dst, ctype, args...);
+      break;
+    case float32:
+      copy<float>(src, dst, ctype, args...);
+      break;
+    case bfloat16:
+      copy<bfloat16_t>(src, dst, ctype, args...);
+      break;
+    case complex64:
+      copy<complex64_t>(src, dst, ctype, args...);
       break;
   }
 }
@@ -242,47 +385,7 @@ void copy(const array& src, array& dst, CopyType ctype) {
 } // namespace
 
 void copy_inplace(const array& src, array& dst, CopyType ctype) {
-  switch (src.dtype()) {
-    case bool_:
-      copy<bool>(src, dst, ctype);
-      break;
-    case uint8:
-      copy<uint8_t>(src, dst, ctype);
-      break;
-    case uint16:
-      copy<uint16_t>(src, dst, ctype);
-      break;
-    case uint32:
-      copy<uint32_t>(src, dst, ctype);
-      break;
-    case uint64:
-      copy<uint64_t>(src, dst, ctype);
-      break;
-    case int8:
-      copy<int8_t>(src, dst, ctype);
-      break;
-    case int16:
-      copy<int16_t>(src, dst, ctype);
-      break;
-    case int32:
-      copy<int32_t>(src, dst, ctype);
-      break;
-    case int64:
-      copy<int64_t>(src, dst, ctype);
-      break;
-    case float16:
-      copy<float16_t>(src, dst, ctype);
-      break;
-    case float32:
-      copy<float>(src, dst, ctype);
-      break;
-    case bfloat16:
-      copy<bfloat16_t>(src, dst, ctype);
-      break;
-    case complex64:
-      copy<complex64_t>(src, dst, ctype);
-      break;
-  }
+  return copy_inplace_dispatch(src, dst, ctype);
 }
 
 void copy(const array& src, array& dst, CopyType ctype) {
@@ -312,4 +415,62 @@ void copy(const array& src, array& dst, CopyType ctype) {
   copy_inplace(src, dst, ctype);
 }
 
+template <typename stride_t>
+void copy_inplace(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype) {
+  switch (ctype) {
+    case CopyType::General:
+    case CopyType::GeneralGeneral:
+      return copy_inplace_dispatch(
+          src,
+          dst,
+          ctype,
+          data_shape,
+          i_strides,
+          o_strides,
+          i_offset,
+          o_offset);
+
+    case CopyType::Scalar:
+    case CopyType::Vector:
+      return copy_inplace_dispatch(src, dst, ctype);
+  }
+}
+
+template <>
+void copy_inplace<int64_t>(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<int64_t>& i_strides,
+    const std::vector<int64_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype) {
+  switch (ctype) {
+    case CopyType::General:
+    case CopyType::GeneralGeneral:
+      return copy_inplace_dispatch(
+          src,
+          dst,
+          ctype,
+          data_shape,
+          i_strides,
+          o_strides,
+          i_offset,
+          o_offset);
+
+    case CopyType::Scalar:
+    case CopyType::Vector:
+      return copy_inplace_dispatch(src, dst, ctype);
+  }
+}
+
 } // namespace mlx::core

mlx/backend/common/copy.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -26,4 +26,15 @@ enum class CopyType {
 void copy(const array& src, array& dst, CopyType ctype);
 void copy_inplace(const array& src, array& dst, CopyType ctype);
 
+template <typename stride_t>
+void copy_inplace(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype);
+
 } // namespace mlx::core

mlx/backend/common/default_primitives.cpp

@@ -94,6 +94,7 @@ DEFAULT(Sign)
 DEFAULT(Sin)
 DEFAULT(Sinh)
 DEFAULT(Slice)
+DEFAULT(SliceUpdate)
 DEFAULT(Softmax)
 DEFAULT(Sort)
 DEFAULT_MULTI(Split)
@@ -105,6 +106,7 @@ DEFAULT_MULTI(SVD)
 DEFAULT(Tan)
 DEFAULT(Tanh)
 DEFAULT(Transpose)
+DEFAULT(Inverse)
 
 namespace {
 

mlx/backend/common/inverse.cpp

@@ -0,0 +1,104 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/copy.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+#ifdef ACCELERATE_NEW_LAPACK
+#include <Accelerate/Accelerate.h>
+#else
+#include <lapack.h>
+#endif
+
+namespace mlx::core {
+
+void inverse_impl(const array& a, array& inv) {
+  // Lapack uses the column-major convention. We take advantage of the following
+  // identity to avoid transposing (see
+  // https://math.stackexchange.com/a/340234):
+  //   (A⁻¹)ᵀ = (Aᵀ)⁻¹
+
+  // The inverse is computed in place, so just copy the input to the output.
+  copy(a, inv, a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+
+  const int N = a.shape(-1);
+  const size_t num_matrices = a.size() / (N * N);
+
+  int info;
+  auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
+
+  for (int i = 0; i < num_matrices; i++) {
+    // Compute LU factorization.
+    sgetrf_(
+        /* m = */ &N,
+        /* n = */ &N,
+        /* a = */ inv.data<float>() + N * N * i,
+        /* lda = */ &N,
+        /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "inverse_impl: LU factorization failed with error code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    static const int lwork_query = -1;
+    float workspace_size = 0;
+
+    // Compute workspace size.
+    sgetri_(
+        /* m = */ &N,
+        /* a = */ nullptr,
+        /* lda = */ &N,
+        /* ipiv = */ nullptr,
+        /* work = */ &workspace_size,
+        /* lwork = */ &lwork_query,
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "inverse_impl: LU workspace calculation failed with error code "
+         << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    const int lwork = workspace_size;
+    auto scratch =
+        array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+
+    // Compute inverse.
+    sgetri_(
+        /* m = */ &N,
+        /* a = */ inv.data<float>() + N * N * i,
+        /* lda = */ &N,
+        /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+        /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "inverse_impl: inversion failed with error code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+}
+
+void Inverse::eval(const std::vector<array>& inputs, array& output) {
+  if (inputs[0].dtype() != float32) {
+    throw std::runtime_error("[Inverse::eval] only supports float32.");
+  }
+  inverse_impl(inputs[0], output);
+}
+
+std::pair<std::vector<array>, std::vector<int>> Inverse::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  return {{linalg::inv(a, stream())}, {ax}};
+}
+
+} // namespace mlx::core

mlx/backend/common/primitives.cpp

@@ -651,36 +651,33 @@ void Sinh::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void Slice::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  if (out.size() == 0) {
-    out.set_data(nullptr);
-    return;
-  }
-  auto& in = inputs[0];
-  auto strides = in.strides();
-  auto flags = in.flags();
-  size_t data_offset = 0;
+std::tuple<bool, int64_t, std::vector<int64_t>> Slice::prepare_slice(
+    const array& in) {
+  int64_t data_offset = 0;
+  bool copy_needed = false;
+  std::vector<int64_t> inp_strides(in.ndim(), 0);
   for (int i = 0; i < in.ndim(); ++i) {
     data_offset += start_indices_[i] * in.strides()[i];
-    strides[i] *= strides_[i];
+    inp_strides[i] = in.strides()[i] * strides_[i];
+
+    copy_needed |= strides_[i] < 0;
   }
 
+  return std::make_tuple(copy_needed, data_offset, inp_strides);
+}
+
+void Slice::shared_buffer_slice(
+    const array& in,
+    const std::vector<size_t>& out_strides,
+    size_t data_offset,
+    array& out) {
   // Compute row/col contiguity
-  size_t data_size = 1;
-  size_t f_stride = 1;
-  size_t b_stride = 1;
-  flags.row_contiguous = true;
-  flags.col_contiguous = true;
-  for (int i = 0, ri = out.ndim() - 1; ri >= 0; i++, ri--) {
-    flags.col_contiguous &= strides[i] == f_stride || out.shape(i) == 1;
-    flags.row_contiguous &= strides[ri] == b_stride || out.shape(ri) == 1;
-    f_stride *= out.shape(i);
-    b_stride *= out.shape(ri);
-    if (strides[i] > 0) {
-      data_size *= out.shape(i);
-    }
-  }
+  auto [data_size, is_row_contiguous, is_col_contiguous] =
+      check_contiguity(out.shape(), out_strides);
+
+  auto flags = in.flags();
+  flags.row_contiguous = is_row_contiguous;
+  flags.col_contiguous = is_col_contiguous;
 
   if (data_size == 1) {
     // Broadcasted scalar array is contiguous.
@@ -694,7 +691,87 @@ void Slice::eval(const std::vector<array>& inputs, array& out) {
     flags.contiguous &= flags.row_contiguous || flags.col_contiguous;
   }
 
-  out.copy_shared_buffer(in, strides, flags, data_size, data_offset);
+  out.copy_shared_buffer(in, out_strides, flags, data_size, data_offset);
+}
+
+void Slice::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+
+  auto& in = inputs[0];
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [copy_needed, data_offset, inp_strides] = prepare_slice(in);
+
+  // Do copy if needed
+  if (copy_needed) {
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    std::vector<int64_t> ostrides{out.strides().begin(), out.strides().end()};
+    copy_inplace<int64_t>(
+        /* const array& src = */ in,
+        /* array& dst = */ out,
+        /* const std::vector<int>& data_shape = */ out.shape(),
+        /* const std::vector<stride_t>& i_strides = */ inp_strides,
+        /* const std::vector<stride_t>& o_strides = */ ostrides,
+        /* int64_t i_offset = */ data_offset,
+        /* int64_t o_offset = */ 0,
+        /* CopyType ctype = */ CopyType::General);
+  } else {
+    std::vector<size_t> ostrides{inp_strides.begin(), inp_strides.end()};
+    shared_buffer_slice(in, ostrides, data_offset, out);
+  }
+}
+
+std::tuple<int64_t, std::vector<int64_t>> SliceUpdate::prepare_slice(
+    const array& in) {
+  int64_t data_offset = 0;
+  std::vector<int64_t> inp_strides(in.ndim(), 0);
+  for (int i = 0; i < in.ndim(); ++i) {
+    data_offset += start_indices_[i] * in.strides()[i];
+    inp_strides[i] = in.strides()[i] * strides_[i];
+  }
+
+  return std::make_tuple(data_offset, inp_strides);
+}
+
+void SliceUpdate::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype);
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] = prepare_slice(out);
+
+  // Do copy
+  std::vector<int64_t> upd_strides{upd.strides().begin(), upd.strides().end()};
+  copy_inplace<int64_t>(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd_strides,
+      /* const std::vector<stride_t>& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral);
 }
 
 void Split::eval(

mlx/backend/common/rope.cpp

@@ -1,13 +0,0 @@
-// Copyright © 2023-2024 Apple Inc.
-
-#include "mlx/fast_primitives.h"
-
-namespace mlx::core::fast {
-
-void RoPE::eval_cpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  throw std::runtime_error("NYI");
-}
-
-} // namespace mlx::core::fast

mlx/backend/common/softmax.cpp

@@ -67,11 +67,15 @@ void Softmax::eval(const std::vector<array>& inputs, array& out) {
     }
   };
   array in = check_input(std::move(inputs[0]));
-  out.set_data(
-      allocator::malloc_or_wait(in.data_size() * in.itemsize()),
-      in.data_size(),
-      in.strides(),
-      in.flags());
+  if (in.is_donatable()) {
+    out.copy_shared_buffer(in);
+  } else {
+    out.set_data(
+        allocator::malloc_or_wait(in.data_size() * in.itemsize()),
+        in.data_size(),
+        in.strides(),
+        in.flags());
+  }
 
   switch (in.dtype()) {
     case bool_:

mlx/backend/common/svd.cpp

@@ -3,6 +3,7 @@
 #include "mlx/allocator.h"
 #include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/lapack_helper.h"
+#include "mlx/linalg.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
@@ -49,8 +50,7 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
 
   // Will contain the indices of eigenvectors that failed to converge (not used
   // here but required by lapack).
-  std::vector<int> iwork;
-  iwork.resize(12 * K);
+  auto iwork = array::Data{allocator::malloc_or_wait(sizeof(int) * 12 * K)};
 
   static const int lwork_query = -1;
 
@@ -82,7 +82,7 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
       /* ldvt = */ &ldvt,
       /* work = */ &workspace_dimension,
       /* lwork = */ &lwork_query,
-      /* iwork = */ iwork.data(),
+      /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
       /* info = */ &info);
 
   if (info != 0) {
@@ -120,7 +120,7 @@ void svd_impl(const array& a, array& u, array& s, array& vt) {
         /* ldvt = */ &ldvt,
         /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
         /* lwork = */ &lwork,
-        /* iwork = */ iwork.data(),
+        /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
         /* info = */ &info);
 
     if (info != 0) {
@@ -145,4 +145,12 @@ void SVD::eval(const std::vector<array>& inputs, std::vector<array>& outputs) {
   svd_impl(inputs[0], outputs[0], outputs[1], outputs[2]);
 }
 
+std::pair<std::vector<array>, std::vector<int>> SVD::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  return {{linalg::svd(a, stream())}, {ax, ax, ax}};
+}
+
 } // namespace mlx::core

mlx/backend/common/utils.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -8,11 +8,12 @@
 
 namespace mlx::core {
 
-inline size_t elem_to_loc(
+template <typename stride_t>
+inline stride_t elem_to_loc(
     int elem,
     const std::vector<int>& shape,
-    const std::vector<size_t>& strides) {
-  size_t loc = 0;
+    const std::vector<stride_t>& strides) {
+  stride_t loc = 0;
   for (int i = shape.size() - 1; i >= 0; --i) {
     auto q_and_r = ldiv(elem, shape[i]);
     loc += q_and_r.rem * strides[i];
@@ -34,10 +35,11 @@ inline size_t elem_to_loc(int elem, const array& a) {
 //
 // When multiple arrays are passed they should all have the same shape. The
 // collapsed axes are also the same so one shape is returned.
-inline std::tuple<std::vector<int>, std::vector<std::vector<size_t>>>
+template <typename stride_t>
+inline std::tuple<std::vector<int>, std::vector<std::vector<stride_t>>>
 collapse_contiguous_dims(
     const std::vector<int>& shape,
-    const std::vector<std::vector<size_t>> strides) {
+    const std::vector<std::vector<stride_t>> strides) {
   // Make a vector that has axes separated with -1. Collapse all axes between
   // -1.
   std::vector<int> to_collapse;
@@ -45,7 +47,7 @@ collapse_contiguous_dims(
     to_collapse.push_back(0);
     for (int i = 1; i < shape.size(); i++) {
       bool contiguous = true;
-      for (const std::vector<size_t>& st : strides) {
+      for (const std::vector<stride_t>& st : strides) {
         if (st[i] * shape[i] != st[i - 1]) {
           contiguous = false;
         }
@@ -62,7 +64,7 @@ collapse_contiguous_dims(
   }
 
   std::vector<int> out_shape;
-  std::vector<std::vector<size_t>> out_strides(strides.size());
+  std::vector<std::vector<stride_t>> out_strides(strides.size());
   for (int i = 0; i < to_collapse.size(); i++) {
     int current_shape = shape[to_collapse[i]];
     while (to_collapse[++i] != -1) {
@@ -70,7 +72,7 @@ collapse_contiguous_dims(
     }
     out_shape.push_back(current_shape);
     for (int j = 0; j < strides.size(); j++) {
-      const std::vector<size_t>& st = strides[j];
+      const std::vector<stride_t>& st = strides[j];
       out_strides[j].push_back(st[to_collapse[i - 1]]);
     }
   }
@@ -94,4 +96,27 @@ collapse_contiguous_dims(Arrays... xs) {
       std::vector<array>{std::forward<Arrays>(xs)...});
 }
 
+template <typename stride_t>
+inline auto check_contiguity(
+    const std::vector<int>& shape,
+    const std::vector<stride_t>& strides) {
+  size_t data_size = 1;
+  size_t f_stride = 1;
+  size_t b_stride = 1;
+  bool is_row_contiguous = true;
+  bool is_col_contiguous = true;
+
+  for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) {
+    is_row_contiguous &= strides[i] == f_stride || shape[i] == 1;
+    is_col_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
+    f_stride *= shape[i];
+    b_stride *= shape[ri];
+    if (strides[i] > 0) {
+      data_size *= shape[i];
+    }
+  }
+
+  return std::make_tuple(data_size, is_row_contiguous, is_col_contiguous);
+}
+
 } // namespace mlx::core

mlx/backend/metal/CMakeLists.txt

@@ -33,6 +33,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/normalization.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp

mlx/backend/metal/copy.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <sstream>
 
@@ -37,15 +37,22 @@ void copy_gpu(const array& in, array& out, CopyType ctype) {
   copy_gpu(in, out, ctype, out.primitive().stream());
 }
 
+template <typename stride_t>
 void copy_gpu_inplace(
     const array& in,
     array& out,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& strides_in_pre,
+    const std::vector<stride_t>& strides_out_pre,
+    int64_t inp_offset,
+    int64_t out_offset,
     CopyType ctype,
     const Stream& s) {
   // Try to collapse contiguous dims
-  auto [shape, strides] = collapse_contiguous_dims(in, out);
-  auto& strides_in = strides[0];
-  auto& strides_out = strides[1];
+  auto [shape, strides] = collapse_contiguous_dims(
+      data_shape, std::vector{strides_in_pre, strides_out_pre});
+  auto& strides_in_ = strides[0];
+  auto& strides_out_ = strides[1];
 
   auto& d = metal::device(s.device);
   std::ostringstream kname;
@@ -72,39 +79,44 @@ void copy_gpu_inplace(
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
   bool donate_in = in.data_shared_ptr() == nullptr;
-  set_array_buffer(compute_encoder, donate_in ? out : in, 0);
-  set_array_buffer(compute_encoder, out, 1);
+
+  inp_offset *= size_of(in.dtype());
+  out_offset *= size_of(out.dtype());
+
+  set_array_buffer(compute_encoder, donate_in ? out : in, inp_offset, 0);
+  set_array_buffer(compute_encoder, out, out_offset, 1);
 
   if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
-    size_t ndim = shape.size();
+    int ndim = shape.size();
+    std::vector<int64_t> strides_in{strides_in_.begin(), strides_in_.end()};
+    std::vector<int64_t> strides_out{strides_out_.begin(), strides_out_.end()};
+
     if (ndim > 3) {
-      compute_encoder->setBytes(shape.data(), ndim * sizeof(int), 2);
-      compute_encoder->setBytes(strides_in.data(), ndim * sizeof(size_t), 3);
-      if (ctype == CopyType::GeneralGeneral) {
-        compute_encoder->setBytes(strides_out.data(), ndim * sizeof(size_t), 4);
-      }
-    } else {
-      // The shape is implicit in the grid for <= 3D
-      compute_encoder->setBytes(strides_in.data(), ndim * sizeof(size_t), 2);
-      if (ctype == CopyType::GeneralGeneral) {
-        compute_encoder->setBytes(strides_out.data(), ndim * sizeof(size_t), 3);
-      }
+      set_vector_bytes(compute_encoder, shape, ndim, 2);
+    }
+    set_vector_bytes(compute_encoder, strides_in, ndim, 3);
+    if (ctype == CopyType::GeneralGeneral) {
+      set_vector_bytes(compute_encoder, strides_out, ndim, 4);
     }
 
     if (ndim > MAX_BINARY_SPECIALIZED_DIMS) {
-      compute_encoder->setBytes(
-          &ndim, sizeof(int), (ctype == CopyType::GeneralGeneral) ? 5 : 4);
+      compute_encoder->setBytes(&ndim, sizeof(int), 5);
     }
 
     int dim0 = ndim > 0 ? shape[ndim - 1] : 1;
     int dim1 = ndim > 1 ? shape[ndim - 2] : 1;
-    int rest = in.size() / (dim0 * dim1);
+
+    size_t data_size = 1;
+    for (auto& s : shape)
+      data_size *= s;
+    int rest = data_size / (dim0 * dim1);
 
     // NB assuming thread_group_size is a power of 2 larger than 32 x 32
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size != 1024) {
       throw std::runtime_error("[Metal::copy] Must use 1024 sized block");
     }
+
     auto group_dims = get_block_dims(dim0, dim1, rest);
     MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
     compute_encoder->dispatchThreads(grid_dims, group_dims);
@@ -120,4 +132,25 @@ void copy_gpu_inplace(
   }
 }
 
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    CopyType ctype,
+    const Stream& s) {
+  return copy_gpu_inplace(
+      in, out, in.shape(), in.strides(), out.strides(), 0, 0, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const std::vector<int64_t>& istride,
+    int64_t ioffset,
+    CopyType ctype,
+    const Stream& s) {
+  std::vector<int64_t> ostrides{out.strides().begin(), out.strides().end()};
+  return copy_gpu_inplace(
+      in, out, in.shape(), istride, ostrides, ioffset, 0, ctype, s);
+}
+
 } // namespace mlx::core

mlx/backend/metal/copy.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -7,12 +7,34 @@
 
 namespace mlx::core {
 
+// Generic copy inplace
+template <typename stride_t>
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const std::vector<int>& data_shape,
+    const std::vector<stride_t>& i_strides,
+    const std::vector<stride_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype,
+    const Stream& s);
+
 void copy_gpu(const array& src, array& out, CopyType ctype, const Stream& s);
 void copy_gpu(const array& src, array& out, CopyType ctype);
+
 void copy_gpu_inplace(
     const array& src,
     array& out,
     CopyType ctype,
     const Stream& s);
 
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const std::vector<int64_t>& istride,
+    int64_t ioffset,
+    CopyType ctype,
+    const Stream& s);
+
 } // namespace mlx::core

mlx/backend/metal/device.cpp

@@ -20,9 +20,9 @@ namespace mlx::core::metal {
 namespace {
 
 // TODO nicer way to set this or possibly expose as an environment variable
-static constexpr int MAX_BUFFERS_PER_QUEUE = 12;
+constexpr int MAX_BUFFERS_PER_QUEUE = 12;
 
-static constexpr const char* default_mtllib_path = METAL_PATH;
+constexpr const char* default_mtllib_path = METAL_PATH;
 
 auto load_device() {
   auto devices = MTL::CopyAllDevices();

mlx/backend/metal/indexing.cpp

@@ -16,7 +16,7 @@ namespace mlx::core {
 
 namespace {
 
-static constexpr int METAL_MAX_INDEX_ARRAYS = 10;
+constexpr int METAL_MAX_INDEX_ARRAYS = 10;
 
 } // namespace
 

mlx/backend/metal/kernels/CMakeLists.txt

@@ -23,6 +23,8 @@ set(
   "gemv"
   "quantized"
   "random"
+  "rms_norm"
+  "layer_norm"
   "rope"
   "scan"
   "scaled_dot_product_attention"

mlx/backend/metal/kernels/copy.metal

@@ -1,29 +1,29 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/metal/kernels/bf16.h"
 #include "mlx/backend/metal/kernels/utils.h"
 
 template <typename T, typename U>
 [[kernel]] void copy_s(
-    device const T* src,
-    device U* dst,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
     uint index [[thread_position_in_grid]]) {
   dst[index] = static_cast<U>(src[0]);
 }
 
 template <typename T, typename U>
 [[kernel]] void copy_v(
-    device const T* src,
-    device U* dst,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
     uint index [[thread_position_in_grid]]) {
   dst[index] = static_cast<U>(src[index]);
 }
 
 template <typename T, typename U>
 [[kernel]] void copy_g_nd1(
-    device const T* src,
-    device U* dst,
-    constant const size_t& src_stride,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t& src_stride [[buffer(3)]],
     uint index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc_1(index, src_stride);
   dst[index] = static_cast<U>(src[src_idx]);
@@ -31,61 +31,61 @@ template <typename T, typename U>
 
 template <typename T, typename U>
 [[kernel]] void copy_g_nd2(
-    device const T* src,
-    device U* dst,
-    constant const size_t src_strides[2],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
     uint2 index [[thread_position_in_grid]],
     uint2 grid_dim [[threads_per_grid]]) {
   auto src_idx = elem_to_loc_2(index, src_strides);
-  size_t dst_idx = index.x + (size_t)grid_dim.x * index.y;
+  int64_t dst_idx = index.x + (int64_t)grid_dim.x * index.y;
   dst[dst_idx] = static_cast<U>(src[src_idx]);
 }
 
 template <typename T, typename U>
 [[kernel]] void copy_g_nd3(
-    device const T* src,
-    device U* dst,
-    constant const size_t src_strides[3],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
     uint3 index [[thread_position_in_grid]],
     uint3 grid_dim [[threads_per_grid]]) {
   auto src_idx = elem_to_loc_3(index, src_strides);
-  size_t dst_idx = index.x + (size_t)grid_dim.x * (index.y + (size_t)grid_dim.y * index.z);
+  int64_t dst_idx = index.x + (int64_t)grid_dim.x * (index.y + (int64_t)grid_dim.y * index.z);
   dst[dst_idx] = static_cast<U>(src[src_idx]);
 }
 
 template <typename T, typename U, int DIM>
 [[kernel]] void copy_g_nd(
-    device const T* src,
-    device U* dst,
-    constant const int src_shape[DIM],
-    constant const size_t src_strides[DIM],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
     uint3 index [[thread_position_in_grid]],
     uint3 grid_dim [[threads_per_grid]]) {
   auto src_idx = elem_to_loc_nd<DIM>(index, src_shape, src_strides);
-  size_t dst_idx = index.x + (size_t)grid_dim.x * (index.y + (size_t)grid_dim.y * index.z);
+  int64_t dst_idx = index.x + (int64_t)grid_dim.x * (index.y + (int64_t)grid_dim.y * index.z);
   dst[dst_idx] = static_cast<U>(src[src_idx]);
 }
 
 template <typename T, typename U>
 [[kernel]] void copy_g(
-    device const T* src,
-    device U* dst,
-    constant const int* src_shape,
-    constant const size_t* src_strides,
-    constant const int& ndim,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int& ndim [[buffer(5)]],
     uint3 index [[thread_position_in_grid]],
     uint3 grid_dim [[threads_per_grid]]) {
   auto src_idx = elem_to_loc(index, src_shape, src_strides, ndim);
-  size_t dst_idx = index.x + (size_t)grid_dim.x * (index.y + (size_t)grid_dim.y * index.z);
+  int64_t dst_idx = index.x + (int64_t)grid_dim.x * (index.y + (int64_t)grid_dim.y * index.z);
   dst[dst_idx] = static_cast<U>(src[src_idx]);
 }
 
 template <typename T, typename U>
 [[kernel]] void copy_gg_nd1(
-    device const T* src,
-    device U* dst,
-    constant const size_t& src_stride,
-    constant const size_t& dst_stride,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t& src_stride [[buffer(3)]],
+    constant const int64_t& dst_stride [[buffer(4)]],
     uint index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc_1(index, src_stride);
   auto dst_idx = elem_to_loc_1(index, dst_stride);
@@ -94,10 +94,10 @@ template <typename T, typename U>
 
 template <typename T, typename U>
 [[kernel]] void copy_gg_nd2(
-    device const T* src,
-    device U* dst,
-    constant const size_t src_strides[2],
-    constant const size_t dst_strides[2],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
     uint2 index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc_2(index, src_strides);
   auto dst_idx = elem_to_loc_2(index, dst_strides);
@@ -106,10 +106,10 @@ template <typename T, typename U>
 
 template <typename T, typename U>
 [[kernel]] void copy_gg_nd3(
-    device const T* src,
-    device U* dst,
-    constant const size_t src_strides[3],
-    constant const size_t dst_strides[3],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
     uint3 index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc_3(index, src_strides);
   auto dst_idx = elem_to_loc_3(index, dst_strides);
@@ -118,11 +118,11 @@ template <typename T, typename U>
 
 template <typename T, typename U, int DIM>
 [[kernel]] void copy_gg_nd(
-    device const T* src,
-    device U* dst,
-    constant const int src_shape[DIM],
-    constant const size_t src_strides[DIM],
-    constant const size_t dst_strides[DIM],
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
     uint3 index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc_nd<DIM>(index, src_shape, src_strides);
   auto dst_idx = elem_to_loc_nd<DIM>(index, src_shape, dst_strides);
@@ -131,12 +131,12 @@ template <typename T, typename U, int DIM>
 
 template <typename T, typename U>
 [[kernel]] void copy_gg(
-    device const T* src,
-    device U* dst,
-    constant const int* src_shape,
-    constant const size_t* src_strides,
-    constant const size_t* dst_strides,
-    constant const int& ndim,
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    constant const int& ndim [[buffer(5)]],
     uint3 index [[thread_position_in_grid]]) {
   auto src_idx = elem_to_loc(index, src_shape, src_strides, ndim);
   auto dst_idx = elem_to_loc(index, src_shape, dst_strides, ndim);
@@ -146,70 +146,70 @@ template <typename T, typename U>
 #define instantiate_copy(name, itype, otype, ctype) \
   template [[host_name(name)]] \
   [[kernel]] void copy_##ctype<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
       uint index [[thread_position_in_grid]]);
 
 #define instantiate_copy_g_dim(name, itype, otype, dims) \
   template [[host_name(name "_" #dims)]] \
   [[kernel]] void copy_g_nd<itype, otype, dims>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const int src_shape[dims], \
-      constant const size_t src_strides[dims], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int* src_shape [[buffer(2)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
       uint3 index [[thread_position_in_grid]], \
       uint3 grid_dim [[threads_per_grid]]); \
   template [[host_name("g" name "_" #dims)]] \
   [[kernel]] void copy_gg_nd<itype, otype, dims>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const int src_shape[dims], \
-      constant const size_t src_strides[dims], \
-      constant const size_t dst_strides[dims], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int* src_shape [[buffer(2)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
+      constant const int64_t* dst_strides [[buffer(4)]], \
       uint3 index [[thread_position_in_grid]]);
 
 
 #define instantiate_copy_g_nd(name, itype, otype) \
   template [[host_name(name "_1")]] \
   [[kernel]] void copy_g_nd1<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t& src_stride, \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t& src_stride [[buffer(3)]], \
       uint index [[thread_position_in_grid]]); \
   template [[host_name(name "_2")]] \
   [[kernel]] void copy_g_nd2<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t src_strides[2], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
       uint2 index [[thread_position_in_grid]], \
       uint2 grid_dim [[threads_per_grid]]); \
   template [[host_name(name "_3")]] \
   [[kernel]] void copy_g_nd3<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t src_strides[3], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
       uint3 index [[thread_position_in_grid]], \
       uint3 grid_dim [[threads_per_grid]]); \
   template [[host_name("g" name "_1")]] \
   [[kernel]] void copy_gg_nd1<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t& src_stride, \
-      constant const size_t& dst_stride, \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t& src_stride [[buffer(3)]], \
+      constant const int64_t& dst_stride [[buffer(4)]], \
       uint index [[thread_position_in_grid]]); \
   template [[host_name("g" name "_2")]] \
   [[kernel]] void copy_gg_nd2<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t src_strides[2], \
-      constant const size_t dst_strides[2], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
+      constant const int64_t* dst_strides [[buffer(4)]], \
       uint2 index [[thread_position_in_grid]]); \
   template [[host_name("g" name "_3")]] \
   [[kernel]] void copy_gg_nd3<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const size_t src_strides[3], \
-      constant const size_t dst_strides[3], \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
+      constant const int64_t* dst_strides [[buffer(4)]], \
       uint3 index [[thread_position_in_grid]]); \
   instantiate_copy_g_dim(name, itype, otype, 4) \
   instantiate_copy_g_dim(name, itype, otype, 5)
@@ -218,21 +218,21 @@ template <typename T, typename U>
 #define instantiate_copy_g(name, itype, otype) \
   template [[host_name(name)]] \
   [[kernel]] void copy_g<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const int* src_shape, \
-      constant const size_t* src_strides, \
-      constant const int& ndim, \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int* src_shape [[buffer(2)]], \
+      constant const int64_t* src_strides  [[buffer(3)]], \
+      constant const int& ndim [[buffer(5)]], \
       uint3 index [[thread_position_in_grid]], \
       uint3 grid_dim [[threads_per_grid]]); \
   template [[host_name("g" name)]] \
   [[kernel]] void copy_gg<itype, otype>( \
-      device const itype* src, \
-      device otype* dst, \
-      constant const int* src_shape, \
-      constant const size_t* src_strides, \
-      constant const size_t* dst_strides, \
-      constant const int& ndim, \
+      device const itype* src [[buffer(0)]], \
+      device otype* dst [[buffer(1)]], \
+      constant const int* src_shape [[buffer(2)]], \
+      constant const int64_t* src_strides [[buffer(3)]], \
+      constant const int64_t* dst_strides [[buffer(4)]], \
+      constant const int& ndim [[buffer(5)]], \
       uint3 index [[thread_position_in_grid]]);
 
 #define instantiate_copy_all(tname, itype, otype) \

mlx/backend/metal/kernels/defines.h

@@ -14,3 +14,5 @@ static MTL_CONST constexpr int MAX_REDUCE_SPECIALIZED_DIMS = 4;
 static MTL_CONST constexpr int REDUCE_N_READS = 16;
 static MTL_CONST constexpr int SOFTMAX_N_READS = 4;
 static MTL_CONST constexpr int SOFTMAX_LOOPED_LIMIT = 4096;
+static MTL_CONST constexpr int RMS_N_READS = 4;
+static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096;

mlx/backend/metal/kernels/layer_norm.metal

@@ -0,0 +1,251 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void layer_norm_single_row(
+    const device T* x,
+    const device T* w,
+    const device T* b,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    constant uint& b_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  float sumx = 0;
+  float sumx2 = 0;
+  float thread_x[N_READS];
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup float local_sumx[SIMD_SIZE];
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_mean[1];
+  threadgroup float local_normalizer[1];
+
+  x += gid * axis_size + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  b += b_stride * lid * N_READS;
+
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = x[i];
+      sumx2 += thread_x[i] * thread_x[i];
+      sumx += thread_x[i];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        thread_x[i] = x[i];
+        sumx2 += thread_x[i] * thread_x[i];
+        sumx += thread_x[i];
+      }
+    }
+  }
+
+  sumx = simd_sum(sumx);
+  sumx2 = simd_sum(sumx2);
+
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sumx[simd_lane_id] = 0;
+    local_sumx2[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sumx[simd_group_id] = sumx;
+    local_sumx2[simd_group_id] = sumx2;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    sumx = simd_sum(local_sumx[simd_lane_id]);
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      float mean = sumx / axis_size;
+      float variance = sumx2 / axis_size - mean * mean;
+
+      local_mean[0] = mean;
+      local_normalizer[0] = metal::precise::rsqrt(variance + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  float mean = local_mean[0];
+  float normalizer = local_normalizer[0];
+
+  // Write the outputs
+  out += gid * axis_size + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = (thread_x[i] - mean) * normalizer;
+      out[i] = w[w_stride * i] * static_cast<T>(thread_x[i]) + b[b_stride * i];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        thread_x[i] = (thread_x[i] - mean) * normalizer;
+        out[i] = w[w_stride * i] * static_cast<T>(thread_x[i]) + b[b_stride * i];
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void layer_norm_looped(
+    const device T* x,
+    const device T* w,
+    const device T* b,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    constant uint& b_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  float sumx = 0;
+  float sumx2 = 0;
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup float local_sumx[SIMD_SIZE];
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_mean[1];
+  threadgroup float local_normalizer[1];
+
+  x += gid * axis_size + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  b += b_stride * lid * N_READS;
+
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        sumx2 += xi * xi;
+        sumx += xi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          sumx2 += xi * xi;
+          sumx += xi;
+        }
+      }
+    }
+  }
+
+  sumx = simd_sum(sumx);
+  sumx2 = simd_sum(sumx2);
+
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sumx[simd_lane_id] = 0;
+    local_sumx2[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sumx[simd_group_id] = sumx;
+    local_sumx2[simd_group_id] = sumx2;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    sumx = simd_sum(local_sumx[simd_lane_id]);
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      float mean = sumx / axis_size;
+      float variance = sumx2 / axis_size - mean * mean;
+
+      local_mean[0] = mean;
+      local_normalizer[0] = metal::precise::rsqrt(variance + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  float mean = local_mean[0];
+  float normalizer = local_normalizer[0];
+
+  // Write the outputs
+  out += gid * axis_size + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = (x[r + i] - mean) * normalizer;
+        out[r + i] = w[w_stride * (i + r)] * static_cast<T>(xi) + b[b_stride * (i + r)];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = (x[r + i] - mean) * normalizer;
+          out[r + i] = w[w_stride * (i + r)] * static_cast<T>(xi) + b[b_stride * (i + r)];
+        }
+      }
+    }
+  }
+}
+
+
+// clang-format off
+#define instantiate_layer_norm_single_row(name, itype)        \
+  template [[host_name("layer_norm" #name)]] [[kernel]] void  \
+  layer_norm_single_row<itype>(                               \
+      const device itype* x,                                  \
+      const device itype* w,                                  \
+      const device itype* b,                                  \
+      device itype* out,                                      \
+      constant float& eps,                                    \
+      constant uint& axis_size,                               \
+      constant uint& w_stride,                                \
+      constant uint& b_stride,                                \
+      uint gid [[thread_position_in_grid]],                   \
+      uint lid [[thread_position_in_threadgroup]],            \
+      uint simd_lane_id [[thread_index_in_simdgroup]],        \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_layer_norm_looped(name, itype)                   \
+  template [[host_name("layer_norm_looped" #name)]] [[kernel]] void  \
+  layer_norm_looped<itype>(                                                 \
+      const device itype* x,                                         \
+      const device itype* w,                                         \
+      const device itype* b,                                         \
+      device itype* out,                                             \
+      constant float& eps,                                           \
+      constant uint& axis_size,                                      \
+      constant uint& w_stride,                                       \
+      constant uint& b_stride,                                       \
+      uint gid [[thread_position_in_grid]],                          \
+      uint lid [[thread_position_in_threadgroup]],                   \
+      uint lsize [[threads_per_threadgroup]],                        \
+      uint simd_lane_id [[thread_index_in_simdgroup]],               \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_layer_norm(name, itype)      \
+  instantiate_layer_norm_single_row(name, itype) \
+  instantiate_layer_norm_looped(name, itype)
+
+instantiate_layer_norm(float32, float)
+instantiate_layer_norm(float16, half)
+instantiate_layer_norm(bfloat16, bfloat16_t)
+    // clang-format on
+

mlx/backend/metal/kernels/reduction/kernels/reduce_row.metal

@@ -108,15 +108,17 @@ template <typename T, typename U, typename Op>
     const short i_ed = short(reduction_size);
     const short i_jump = reductions_per_thread;
 
-    for(short r = r_st; r < r_ed; r += r_jump) {
+    if(r_st < r_jump) {
+      for(short r = r_st; r < r_ed; r += r_jump) {
 
-      uint in_idx = elem_to_loc(out_idx + r * out_size, shape, strides, ndim);
-      const device T * in_row = in + in_idx;
+        uint in_idx = elem_to_loc(out_idx + r * out_size, shape, strides, ndim);
+        const device T * in_row = in + in_idx;
+
+        for(short i = i_st; i < i_ed; i += i_jump) {
+          total_val = op(static_cast<U>(in_row[i]), total_val);
+        }
 
-      for(short i = i_st; i < i_ed; i += i_jump) {
-        total_val = op(static_cast<U>(in_row[i]), total_val);
       }
-
     }
 
   }

mlx/backend/metal/kernels/rms_norm.metal

@@ -0,0 +1,194 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_single_row(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    threadgroup float* local_inv_mean [[threadgroup(0)]],
+    threadgroup float* local_sums [[threadgroup(1)]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  float acc = 0;
+  x += gid * axis_size + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      float xi = x[i];
+      acc += xi * xi;
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        float xi = x[i];
+        acc += xi * xi;
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * axis_size + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_looped(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    threadgroup float* local_inv_mean [[threadgroup(0)]],
+    threadgroup float* local_sums [[threadgroup(1)]],
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  float acc = 0;
+  x += gid * axis_size + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        acc += xi * xi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          acc += xi * xi;
+        }
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * axis_size + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        out[r + i] = w[w_stride * (i + r)] *
+            static_cast<T>(x[r + i] * local_inv_mean[0]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          out[r + i] = w[w_stride * (i + r)] *
+              static_cast<T>(x[r + i] * local_inv_mean[0]);
+        }
+      }
+    }
+  }
+}
+
+// clang-format off
+#define instantiate_rms_single_row(name, itype)               \
+  template [[host_name("rms" #name)]] [[kernel]] void         \
+  rms_single_row<itype>(                                      \
+      const device itype* x,                                  \
+      const device itype* w,                                  \
+      device itype* out,                                      \
+      constant float& eps,                                    \
+      constant uint& axis_size,                               \
+      constant uint& w_stride,                                \
+      threadgroup float* local_inv_mean [[threadgroup(0)]],   \
+      threadgroup float* local_sums [[threadgroup(1)]],       \
+      uint gid [[thread_position_in_grid]],                   \
+      uint lid [[thread_position_in_threadgroup]],            \
+      uint simd_lane_id [[thread_index_in_simdgroup]],        \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_rms_looped(name, itype)                   \
+  template [[host_name("rms_looped" #name)]] [[kernel]] void  \
+  rms_looped<itype>(                                          \
+      const device itype* x,                                  \
+      const device itype* w,                                  \
+      device itype* out,                                      \
+      constant float& eps,                                    \
+      constant uint& axis_size,                               \
+      constant uint& w_stride,                                \
+      threadgroup float* local_inv_mean [[threadgroup(0)]],   \
+      threadgroup float* local_sums [[threadgroup(1)]],       \
+      uint gid [[thread_position_in_grid]],                   \
+      uint lid [[thread_position_in_threadgroup]],            \
+      uint lsize [[threads_per_threadgroup]],                 \
+      uint simd_lane_id [[thread_index_in_simdgroup]],        \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_rms(name, itype)      \
+  instantiate_rms_single_row(name, itype) \
+  instantiate_rms_looped(name, itype)
+
+instantiate_rms(float32, float)
+instantiate_rms(float16, half)
+instantiate_rms(bfloat16, bfloat16_t)
+    // clang-format on

mlx/backend/metal/kernels/softmax.metal

@@ -1,6 +1,5 @@
 // Copyright © 2023 Apple Inc.
 
-#include <metal_atomic>
 #include <metal_common>
 #include <metal_simdgroup>
 
@@ -224,5 +223,6 @@ template <typename T, int N_READS = SOFTMAX_N_READS>
   instantiate_softmax_single_row(name, itype) \
       instantiate_softmax_looped(name, itype)
 
-instantiate_softmax(float32, float) instantiate_softmax(float16, half)
-    instantiate_softmax(bfloat16, bfloat16_t)
+instantiate_softmax(float32, float)
+instantiate_softmax(float16, half)
+instantiate_softmax(bfloat16, bfloat16_t)

mlx/backend/metal/kernels/utils.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -65,12 +65,18 @@ struct Limits<bool> {
 // Indexing utils
 ///////////////////////////////////////////////////////////////////////////////
 
-inline size_t elem_to_loc(
+#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+
+///////////////////////////////////////////////////////////////////////////////
+// Single Array with generic dims
+
+template <typename stride_t>
+METAL_FUNC stride_t elem_to_loc(
     uint elem,
     device const int* shape,
-    device const size_t* strides,
+    device const stride_t* strides,
     int ndim) {
-  size_t loc = 0;
+  stride_t loc = 0;
   for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
     loc += (elem % shape[i]) * strides[i];
     elem /= shape[i];
@@ -78,12 +84,13 @@ inline size_t elem_to_loc(
   return loc;
 }
 
-inline size_t elem_to_loc(
+template <typename stride_t>
+METAL_FUNC stride_t elem_to_loc(
     uint elem,
     constant const int* shape,
-    constant const size_t* strides,
+    constant const stride_t* strides,
     int ndim) {
-  size_t loc = 0;
+  stride_t loc = 0;
   for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
     loc += (elem % shape[i]) * strides[i];
     elem /= shape[i];
@@ -91,52 +98,59 @@ inline size_t elem_to_loc(
   return loc;
 }
 
-template <int NDIM>
-inline uint3 elem_to_loc_3_nd(
+// Non templated version to handle arbitrary dims
+template <typename stride_t>
+METAL_FUNC stride_t elem_to_loc(
     uint3 elem,
-    constant const int shape[NDIM],
-    constant const size_t a_strides[NDIM],
-    constant const size_t b_strides[NDIM],
-    constant const size_t c_strides[NDIM]) {
-  uint3 loc = {
-      static_cast<uint>(
-          elem.x * a_strides[NDIM - 1] + elem.y * a_strides[NDIM - 2]),
-      static_cast<uint>(
-          elem.x * b_strides[NDIM - 1] + elem.y * b_strides[NDIM - 2]),
-      static_cast<uint>(
-          elem.x * c_strides[NDIM - 1] + elem.y * c_strides[NDIM - 2])};
-  for (int d = NDIM - 3; d >= 0; --d) {
-    uint l = elem.z % shape[d];
-    loc.x += l * a_strides[d];
-    loc.y += l * b_strides[d];
-    loc.z += l * c_strides[d];
+    constant const int* shape,
+    constant const stride_t* strides,
+    int ndim) {
+  stride_t loc = elem.x * strides[ndim - 1] + elem.y * strides[ndim - 2];
+  for (int d = ndim - 3; d >= 0; --d) {
+    loc += (elem.z % shape[d]) * strides[d];
     elem.z /= shape[d];
   }
   return loc;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// Single Array with fixed N dims
+
+template <typename stride_t>
+METAL_FUNC stride_t elem_to_loc_1(uint elem, constant const stride_t& stride) {
+  return elem * stride;
+}
+
+template <typename stride_t>
+METAL_FUNC stride_t
+elem_to_loc_2(uint2 elem, constant const stride_t strides[2]) {
+  return elem.x * strides[1] + elem.y * strides[0];
+}
+
+template <typename stride_t>
+METAL_FUNC stride_t
+elem_to_loc_3(uint3 elem, constant const stride_t strides[3]) {
+  return elem.x * strides[2] + elem.y * strides[1] + elem.z * strides[0];
+}
+
 template <int NDIM>
-inline uint2 elem_to_loc_2_nd(
-    uint3 elem,
-    constant const int shape[NDIM],
-    constant const size_t a_strides[NDIM],
-    constant const size_t b_strides[NDIM]) {
-  uint2 loc = {
-      static_cast<uint>(
-          elem.x * a_strides[NDIM - 1] + elem.y * a_strides[NDIM - 2]),
-      static_cast<uint>(
-          elem.x * b_strides[NDIM - 1] + elem.y * b_strides[NDIM - 2])};
-  for (int d = NDIM - 3; d >= 0; --d) {
-    uint l = elem.z % shape[d];
-    loc.x += l * a_strides[d];
-    loc.y += l * b_strides[d];
-    elem.z /= shape[d];
+METAL_FUNC size_t elem_to_loc_nd(
+    uint elem,
+    device const int* shape,
+    device const size_t* strides) {
+  size_t loc = (elem % shape[NDIM - 1]) * strides[NDIM - 1];
+
+  MLX_MTL_PRAGMA_UNROLL
+  for (int d = NDIM - 2; d >= 0; --d) {
+    elem /= shape[d + 1];
+    loc += (elem % shape[d]) * strides[d];
   }
+
   return loc;
 }
 
 template <int NDIM>
-inline size_t elem_to_loc_nd(
+METAL_FUNC size_t elem_to_loc_nd(
     uint3 elem,
     constant const int shape[NDIM],
     constant const size_t strides[NDIM]) {
@@ -148,33 +162,59 @@ inline size_t elem_to_loc_nd(
   return loc;
 }
 
-inline size_t elem_to_loc_1(uint elem, constant const size_t& stride) {
-  return elem * stride;
+template <int NDIM>
+METAL_FUNC int64_t elem_to_loc_nd(
+    uint elem,
+    constant const int shape[NDIM],
+    constant const int64_t strides[NDIM]) {
+  int64_t loc = (elem % shape[NDIM - 1]) * strides[NDIM - 1];
+
+  MLX_MTL_PRAGMA_UNROLL
+  for (int d = NDIM - 2; d >= 0; --d) {
+    elem /= shape[d + 1];
+    loc += (elem % shape[d]) * strides[d];
+  }
+
+  return loc;
 }
 
-inline size_t elem_to_loc_2(uint2 elem, constant const size_t strides[2]) {
-  return elem.x * strides[1] + elem.y * strides[0];
-}
-
-inline size_t elem_to_loc_3(uint3 elem, constant const size_t strides[3]) {
-  return elem.x * strides[2] + elem.y * strides[1] + elem.z * strides[0];
-}
-
-// Non templated version to handle arbitrary dims
-inline size_t elem_to_loc(
+template <int NDIM>
+METAL_FUNC int64_t elem_to_loc_nd(
     uint3 elem,
-    constant const int* shape,
-    constant const size_t* strides,
-    int ndim) {
-  size_t loc = elem.x * strides[ndim - 1] + elem.y * strides[ndim - 2];
-  for (int d = ndim - 3; d >= 0; --d) {
+    constant const int shape[NDIM],
+    constant const int64_t strides[NDIM]) {
+  int64_t loc = elem.x * strides[NDIM - 1] + elem.y * strides[NDIM - 2];
+  for (int d = NDIM - 3; d >= 0; --d) {
     loc += (elem.z % shape[d]) * strides[d];
     elem.z /= shape[d];
   }
   return loc;
 }
 
-inline uint3 elem_to_loc_3_nd(
+///////////////////////////////////////////////////////////////////////////////
+// Multiple Arrays with generic dims
+
+METAL_FUNC uint2 elem_to_loc_2_nd(
+    uint3 elem,
+    constant const int* shape,
+    constant const size_t* a_strides,
+    constant const size_t* b_strides,
+    int ndim) {
+  uint2 loc = {
+      static_cast<uint>(
+          elem.x * a_strides[ndim - 1] + elem.y * a_strides[ndim - 2]),
+      static_cast<uint>(
+          elem.x * b_strides[ndim - 1] + elem.y * b_strides[ndim - 2])};
+  for (int d = ndim - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * a_strides[d];
+    loc.y += l * b_strides[d];
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
+METAL_FUNC uint3 elem_to_loc_3_nd(
     uint3 elem,
     constant const int* shape,
     constant const size_t* a_strides,
@@ -198,18 +238,21 @@ inline uint3 elem_to_loc_3_nd(
   return loc;
 }
 
-inline uint2 elem_to_loc_2_nd(
+///////////////////////////////////////////////////////////////////////////////
+// Multiple Arrays with fixed N dims
+
+template <int NDIM>
+METAL_FUNC uint2 elem_to_loc_2_nd(
     uint3 elem,
-    constant const int* shape,
-    constant const size_t* a_strides,
-    constant const size_t* b_strides,
-    int ndim) {
+    constant const int shape[NDIM],
+    constant const size_t a_strides[NDIM],
+    constant const size_t b_strides[NDIM]) {
   uint2 loc = {
       static_cast<uint>(
-          elem.x * a_strides[ndim - 1] + elem.y * a_strides[ndim - 2]),
+          elem.x * a_strides[NDIM - 1] + elem.y * a_strides[NDIM - 2]),
       static_cast<uint>(
-          elem.x * b_strides[ndim - 1] + elem.y * b_strides[ndim - 2])};
-  for (int d = ndim - 3; d >= 0; --d) {
+          elem.x * b_strides[NDIM - 1] + elem.y * b_strides[NDIM - 2])};
+  for (int d = NDIM - 3; d >= 0; --d) {
     uint l = elem.z % shape[d];
     loc.x += l * a_strides[d];
     loc.y += l * b_strides[d];
@@ -219,55 +262,26 @@ inline uint2 elem_to_loc_2_nd(
 }
 
 template <int NDIM>
-inline uint elem_to_loc_nd(
-    uint elem,
-    device const int* shape,
-    device const size_t* strides);
-
-template <>
-inline uint elem_to_loc_nd<1>(
-    uint elem,
-    device const int* shape,
-    device const size_t* strides) {
-  return (elem % shape[0]) * strides[0];
-}
-
-template <>
-inline uint elem_to_loc_nd<2>(
-    uint elem,
-    device const int* shape,
-    device const size_t* strides) {
-  uint loc = (elem % shape[1]) * strides[1];
-  elem /= shape[1];
-  loc += (elem % shape[0]) * strides[0];
-  return loc;
-}
-
-template <>
-inline uint elem_to_loc_nd<3>(
-    uint elem,
-    device const int* shape,
-    device const size_t* strides) {
-  uint loc = (elem % shape[2]) * strides[2];
-  elem /= shape[2];
-  loc += (elem % shape[1]) * strides[1];
-  elem /= shape[1];
-  loc += (elem % shape[0]) * strides[0];
-  return loc;
-}
-
-template <>
-inline uint elem_to_loc_nd<4>(
-    uint elem,
-    device const int* shape,
-    device const size_t* strides) {
-  uint loc = (elem % shape[3]) * strides[3];
-  elem /= shape[3];
-  loc += (elem % shape[2]) * strides[2];
-  elem /= shape[2];
-  loc += (elem % shape[1]) * strides[1];
-  elem /= shape[1];
-  loc += (elem % shape[0]) * strides[0];
+METAL_FUNC uint3 elem_to_loc_3_nd(
+    uint3 elem,
+    constant const int shape[NDIM],
+    constant const size_t a_strides[NDIM],
+    constant const size_t b_strides[NDIM],
+    constant const size_t c_strides[NDIM]) {
+  uint3 loc = {
+      static_cast<uint>(
+          elem.x * a_strides[NDIM - 1] + elem.y * a_strides[NDIM - 2]),
+      static_cast<uint>(
+          elem.x * b_strides[NDIM - 1] + elem.y * b_strides[NDIM - 2]),
+      static_cast<uint>(
+          elem.x * c_strides[NDIM - 1] + elem.y * c_strides[NDIM - 2])};
+  for (int d = NDIM - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * a_strides[d];
+    loc.y += l * b_strides[d];
+    loc.z += l * c_strides[d];
+    elem.z /= shape[d];
+  }
   return loc;
 }
 

mlx/backend/metal/matmul.cpp

@@ -206,7 +206,7 @@ inline auto collapse_batches(const array& a, const array& b) {
   std::vector<size_t> B_bstride{b.strides().begin(), b.strides().end() - 2};
 
   auto [batch_shape, batch_strides] =
-      collapse_contiguous_dims(A_bshape, {A_bstride, B_bstride});
+      collapse_contiguous_dims(A_bshape, std::vector{A_bstride, B_bstride});
 
   auto A_batch_stride = batch_strides[0];
   auto B_batch_stride = batch_strides[1];
@@ -237,8 +237,8 @@ inline auto collapse_batches(const array& a, const array& b, const array& c) {
   std::vector<size_t> B_bstride{b.strides().begin(), b.strides().end() - 2};
   std::vector<size_t> C_bstride{c.strides().begin(), c.strides().end() - 2};
 
-  auto [batch_shape, batch_strides] =
-      collapse_contiguous_dims(A_bshape, {A_bstride, B_bstride, C_bstride});
+  auto [batch_shape, batch_strides] = collapse_contiguous_dims(
+      A_bshape, std::vector{A_bstride, B_bstride, C_bstride});
 
   auto A_batch_stride = batch_strides[0];
   auto B_batch_stride = batch_strides[1];

mlx/backend/metal/normalization.cpp

@@ -0,0 +1,185 @@
+// Copyright © 2024 Apple Inc.
+#include <algorithm>
+
+#include "mlx/backend/metal/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+void RMSNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& out = outputs[0];
+
+  // Make sure that the last dimension is contiguous
+  std::vector<array> copies;
+  auto check_input = [&copies, &s](const array& x) {
+    bool no_copy = x.strides()[x.ndim() - 1] == 1;
+    if (x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back());
+    }
+    if (no_copy) {
+      return x;
+    } else {
+      array x_copy(x.shape(), x.dtype(), nullptr, {});
+      copy_gpu(x, x_copy, CopyType::General, s);
+      copies.push_back(x_copy);
+      return x_copy;
+    }
+  };
+  const array& x = check_input(inputs[0]);
+  const array& w = inputs[1];
+
+  if (x.is_donatable()) {
+    out.move_shared_buffer(x);
+  } else {
+    out.set_data(
+        allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+        x.data_size(),
+        x.strides(),
+        x.flags());
+  }
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = RMS_N_READS;
+  const int looped_limit = RMS_LOOPED_LIMIT;
+  std::string op_name = "rms";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+  }
+  op_name += type_to_name(out);
+  auto compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = w.strides()[0];
+    compute_encoder->setComputePipelineState(kernel);
+    set_array_buffer(
+        compute_encoder, x.data_shared_ptr() == nullptr ? out : x, 0);
+    set_array_buffer(compute_encoder, w, 1);
+    set_array_buffer(compute_encoder, out, 2);
+    compute_encoder->setBytes(&eps_, sizeof(float), 3);
+    compute_encoder->setBytes(&axis_size, sizeof(int), 4);
+    compute_encoder->setBytes(&w_stride, sizeof(uint32_t), 5);
+    compute_encoder->setThreadgroupMemoryLength(
+        16 * 8, 0); // minimum of 16 bytes
+    compute_encoder->setThreadgroupMemoryLength(simd_size * sizeof(float), 1);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+  }
+  d.get_command_buffer(s.index)->addCompletedHandler(
+      [copies](MTL::CommandBuffer*) mutable { copies.clear(); });
+}
+
+void LayerNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& out = outputs[0];
+
+  // Make sure that the last dimension is contiguous
+  std::vector<array> copies;
+  auto check_input = [&copies, &s](const array& x) {
+    bool no_copy = x.strides()[x.ndim() - 1] == 1;
+    if (x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back());
+    }
+    if (no_copy) {
+      return x;
+    } else {
+      array x_copy(x.shape(), x.dtype(), nullptr, {});
+      copy_gpu(x, x_copy, CopyType::General, s);
+      copies.push_back(x_copy);
+      return x_copy;
+    }
+  };
+  const array& x = check_input(inputs[0]);
+  const array& w = inputs[1];
+  const array& b = inputs[2];
+
+  if (x.is_donatable()) {
+    out.move_shared_buffer(x);
+  } else {
+    out.set_data(
+        allocator::malloc_or_wait(x.data_size() * x.itemsize()),
+        x.data_size(),
+        x.strides(),
+        x.flags());
+  }
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = RMS_N_READS;
+  const int looped_limit = RMS_LOOPED_LIMIT;
+  std::string op_name = "layer_norm";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+  }
+  op_name += type_to_name(out);
+  auto compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    uint32_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0;
+    compute_encoder->setComputePipelineState(kernel);
+    set_array_buffer(
+        compute_encoder, x.data_shared_ptr() == nullptr ? out : x, 0);
+    set_array_buffer(compute_encoder, w, 1);
+    set_array_buffer(compute_encoder, b, 2);
+    set_array_buffer(compute_encoder, out, 3);
+    compute_encoder->setBytes(&eps_, sizeof(float), 4);
+    compute_encoder->setBytes(&axis_size, sizeof(int), 5);
+    compute_encoder->setBytes(&w_stride, sizeof(uint32_t), 6);
+    compute_encoder->setBytes(&b_stride, sizeof(uint32_t), 7);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+  }
+  d.get_command_buffer(s.index)->addCompletedHandler(
+      [copies](MTL::CommandBuffer*) mutable { copies.clear(); });
+}
+
+} // namespace mlx::core::fast

mlx/backend/metal/primitives.cpp

@@ -17,7 +17,7 @@ namespace mlx::core {
 
 namespace {
 
-static constexpr int METAL_MAX_INDEX_ARRAYS = 10;
+constexpr int METAL_MAX_INDEX_ARRAYS = 10;
 
 void binary_op(
     const std::vector<array>& inputs,
@@ -865,7 +865,73 @@ void Sqrt::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Slice::eval_gpu(const std::vector<array>& inputs, array& out) {
-  eval(inputs, out);
+  assert(inputs.size() == 1);
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+
+  auto& in = inputs[0];
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [copy_needed, data_offset, inp_strides] = prepare_slice(in);
+
+  // Do copy if needed
+  if (copy_needed) {
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    std::vector<int64_t> ostrides{out.strides().begin(), out.strides().end()};
+    copy_gpu_inplace(
+        /* const array& in = */ in,
+        /* array& out = */ out,
+        /* const std::vector<int>& data_shape = */ out.shape(),
+        /* const std::vector<stride_t>& i_strides = */ inp_strides,
+        /* const std::vector<stride_t>& o_strides = */ ostrides,
+        /* int64_t i_offset = */ data_offset,
+        /* int64_t o_offset = */ 0,
+        /* CopyType ctype = */ CopyType::General,
+        /* const Stream& s = */ stream());
+  } else {
+    std::vector<size_t> ostrides{inp_strides.begin(), inp_strides.end()};
+    shared_buffer_slice(in, ostrides, data_offset, out);
+  }
+}
+
+void SliceUpdate::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] = prepare_slice(out);
+
+  // Do copy
+  std::vector<int64_t> upd_strides{upd.strides().begin(), upd.strides().end()};
+  copy_gpu_inplace<int64_t>(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd_strides,
+      /* const std::vector<stride_t>& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      /* const Stream& s = */ stream());
 }
 
 void StopGradient::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -900,4 +966,8 @@ void SVD::eval_gpu(
   throw std::runtime_error("[SVD::eval_gpu] Metal SVD NYI.");
 }
 
+void Inverse::eval_gpu(const std::vector<array>& inputs, array& output) {
+  throw std::runtime_error("[Inverse::eval_gpu] Metal inversion NYI.");
+}
+
 } // namespace mlx::core

mlx/backend/metal/rope.cpp

@@ -1,5 +1,5 @@
 // Copyright © 2023-2024 Apple Inc.
-
+#include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/utils.h"
 #include "mlx/fast_primitives.h"
 
@@ -13,39 +13,63 @@ void RoPE::eval_gpu(
   auto& in = inputs[0];
   auto& out = outputs[0];
 
-  if (in.ndim() != 3) {
-    throw std::runtime_error(
-        "[RoPE] Only 3 dimensions are supported (batch x sequence x dims)");
+  if (in.ndim() < 3) {
+    throw std::runtime_error("[RoPE] Input must have at least 3 dimensions");
   }
   if (dims_ != in.shape(-1)) {
     throw std::runtime_error("[RoPE] Partial RoPE application not supported");
   }
-  if (in.flags().row_contiguous && in.is_donatable()) {
-    out.move_shared_buffer(in);
-  } else {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  }
 
   auto& s = out.primitive().stream();
   auto& d = metal::device(s.device);
+
+  size_t strides[3];
+  bool donated = false;
+  int ndim = in.ndim();
+  size_t mat_size = in.shape()[ndim - 2] * in.shape()[ndim - 1];
+  if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      donated = true;
+      out.move_shared_buffer(in);
+    } else {
+      out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    }
+    strides[0] = mat_size;
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (ndim == 3) {
+    // Handle non-contiguous 3D inputs
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    strides[0] = in.strides()[0];
+    strides[1] = in.strides()[1];
+    strides[2] = in.strides()[2];
+  } else {
+    // Copy non-contiguous > 3D inputs into the output and treat
+    // input as donated
+    donated = true;
+    copy_gpu(in, out, CopyType::General, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  }
+
   std::ostringstream kname;
   kname << "rope_" << (traditional_ ? "traditional_" : "") << type_to_name(in);
   auto kernel = d.get_kernel(kname.str());
   auto compute_encoder = d.get_command_encoder(s.index);
 
-  bool donated = in.data_shared_ptr() == nullptr;
   float base = std::log2(base_);
   compute_encoder->setComputePipelineState(kernel);
   set_array_buffer(compute_encoder, donated ? out : in, 0);
   set_array_buffer(compute_encoder, out, 1);
-  compute_encoder->setBytes(in.strides().data(), 3 * sizeof(size_t), 2);
+  compute_encoder->setBytes(&strides, 3 * sizeof(size_t), 2);
   compute_encoder->setBytes(&offset_, sizeof(int), 3);
   compute_encoder->setBytes(&base, sizeof(float), 4);
   compute_encoder->setBytes(&scale_, sizeof(float), 5);
 
-  int dim0 = in.shape(2) / 2;
-  int dim1 = in.shape(1);
-  int dim2 = in.shape(0);
+  int dim0 = in.shape()[ndim - 1] / 2;
+  int dim1 = in.shape()[ndim - 2];
+  int dim2 = in.size() / mat_size;
   auto group_dims = get_block_dims(dim0, dim1, dim2);
   auto grid_dims = MTL::Size(dim0, dim1, dim2);
   compute_encoder->dispatchThreads(grid_dims, group_dims);

mlx/backend/metal/softmax.cpp

@@ -37,11 +37,15 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
     }
   };
   const array& in = check_input(inputs[0]);
-  out.set_data(
-      allocator::malloc_or_wait(in.data_size() * in.itemsize()),
-      in.data_size(),
-      in.strides(),
-      in.flags());
+  if (in.is_donatable()) {
+    out.move_shared_buffer(in);
+  } else {
+    out.set_data(
+        allocator::malloc_or_wait(in.data_size() * in.itemsize()),
+        in.data_size(),
+        in.strides(),
+        in.flags());
+  }
 
   int axis_size = in.shape().back();
   int n_rows = in.data_size() / axis_size;
@@ -75,6 +79,8 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
     }
 
     compute_encoder->setComputePipelineState(kernel);
+    set_array_buffer(
+        compute_encoder, in.data_shared_ptr() == nullptr ? out : in, 0);
     set_array_buffer(compute_encoder, in, 0);
     set_array_buffer(compute_encoder, out, 1);
     compute_encoder->setBytes(&axis_size, sizeof(int), 2);

mlx/backend/metal/utils.h

@@ -9,16 +9,43 @@ namespace mlx::core {
 
 namespace {
 
-void set_array_buffer(
-    MTL::ComputeCommandEncoder* enc,
-    const array& a,
-    int idx) {
+inline void
+set_array_buffer(MTL::ComputeCommandEncoder* enc, const array& a, int idx) {
   auto a_buf = static_cast<const MTL::Buffer*>(a.buffer().ptr());
   auto offset = a.data<char>() -
       static_cast<char*>(const_cast<MTL::Buffer*>(a_buf)->contents());
   enc->setBuffer(a_buf, offset, idx);
 }
 
+inline void set_array_buffer(
+    MTL::ComputeCommandEncoder* enc,
+    const array& a,
+    int64_t offset,
+    int idx) {
+  auto a_buf = static_cast<const MTL::Buffer*>(a.buffer().ptr());
+  auto base_offset = a.data<char>() -
+      static_cast<char*>(const_cast<MTL::Buffer*>(a_buf)->contents());
+  base_offset += offset;
+  enc->setBuffer(a_buf, base_offset, idx);
+}
+
+template <typename T>
+inline void set_vector_bytes(
+    MTL::ComputeCommandEncoder* enc,
+    const std::vector<T>& vec,
+    size_t nelems,
+    int idx) {
+  enc->setBytes(vec.data(), nelems * sizeof(T), idx);
+}
+
+template <typename T>
+inline void set_vector_bytes(
+    MTL::ComputeCommandEncoder* enc,
+    const std::vector<T>& vec,
+    int idx) {
+  return set_vector_bytes(enc, vec, vec.size(), idx);
+}
+
 std::string type_to_name(const array& a) {
   std::string tname;
   switch (a.dtype()) {

mlx/backend/no_metal/primitives.cpp

@@ -87,6 +87,7 @@ NO_GPU(Sign)
 NO_GPU(Sin)
 NO_GPU(Sinh)
 NO_GPU(Slice)
+NO_GPU(SliceUpdate)
 NO_GPU(Softmax)
 NO_GPU(Sort)
 NO_GPU_MULTI(Split)
@@ -98,8 +99,11 @@ NO_GPU_MULTI(SVD)
 NO_GPU(Tan)
 NO_GPU(Tanh)
 NO_GPU(Transpose)
+NO_GPU(Inverse)
 
 namespace fast {
+NO_GPU_MULTI(LayerNorm)
+NO_GPU_MULTI(RMSNorm)
 NO_GPU_MULTI(RoPE)
 NO_GPU(ScaledDotProductAttention)
 } // namespace fast

mlx/dtype.cpp

@@ -11,9 +11,9 @@ namespace mlx::core {
 
 namespace {
 
-static constexpr int num_types = 13;
+constexpr int num_types = 13;
 
-static constexpr Dtype::Kind type_kinds[num_types] = {
+constexpr Dtype::Kind type_kinds[num_types] = {
     Dtype::Kind::b, // bool_,
     Dtype::Kind::u, // uint8,
     Dtype::Kind::u, // uint16,
@@ -32,7 +32,7 @@ static constexpr Dtype::Kind type_kinds[num_types] = {
 // Following Jax type promotion rules:
 // https://jax.readthedocs.io/en/latest/type_promotion.html
 // clang-format off
-static constexpr Dtype type_rules[num_types][num_types] = {
+constexpr Dtype type_rules[num_types][num_types] = {
 // bool       uint8      uint16     uint32     uint64     int8       int16      int32      int64      float16    float32    bfloat16   complex64
   {bool_,     uint8,     uint16,    uint32,    uint64,    int8,      int16,     int32,     int64,     float16,   float32,   bfloat16,  complex64}, // bool
   {uint8,     uint8,     uint16,    uint32,    uint64,    int16,     int16,     int32,     int64,     float16,   float32,   bfloat16,  complex64}, // uint8

mlx/dtype.h

@@ -46,22 +46,22 @@ struct Dtype {
   };
 };
 
-static constexpr Dtype bool_{Dtype::Val::bool_, sizeof(bool)};
+inline constexpr Dtype bool_{Dtype::Val::bool_, sizeof(bool)};
 
-static constexpr Dtype uint8{Dtype::Val::uint8, sizeof(uint8_t)};
-static constexpr Dtype uint16{Dtype::Val::uint16, sizeof(uint16_t)};
-static constexpr Dtype uint32{Dtype::Val::uint32, sizeof(uint32_t)};
-static constexpr Dtype uint64{Dtype::Val::uint64, sizeof(uint64_t)};
+inline constexpr Dtype uint8{Dtype::Val::uint8, sizeof(uint8_t)};
+inline constexpr Dtype uint16{Dtype::Val::uint16, sizeof(uint16_t)};
+inline constexpr Dtype uint32{Dtype::Val::uint32, sizeof(uint32_t)};
+inline constexpr Dtype uint64{Dtype::Val::uint64, sizeof(uint64_t)};
 
-static constexpr Dtype int8{Dtype::Val::int8, sizeof(int8_t)};
-static constexpr Dtype int16{Dtype::Val::int16, sizeof(int16_t)};
-static constexpr Dtype int32{Dtype::Val::int32, sizeof(int32_t)};
-static constexpr Dtype int64{Dtype::Val::int64, sizeof(int64_t)};
+inline constexpr Dtype int8{Dtype::Val::int8, sizeof(int8_t)};
+inline constexpr Dtype int16{Dtype::Val::int16, sizeof(int16_t)};
+inline constexpr Dtype int32{Dtype::Val::int32, sizeof(int32_t)};
+inline constexpr Dtype int64{Dtype::Val::int64, sizeof(int64_t)};
 
-static constexpr Dtype float16{Dtype::Val::float16, sizeof(uint16_t)};
-static constexpr Dtype float32{Dtype::Val::float32, sizeof(float)};
-static constexpr Dtype bfloat16{Dtype::Val::bfloat16, sizeof(uint16_t)};
-static constexpr Dtype complex64{Dtype::Val::complex64, sizeof(complex64_t)};
+inline constexpr Dtype float16{Dtype::Val::float16, sizeof(uint16_t)};
+inline constexpr Dtype float32{Dtype::Val::float32, sizeof(float)};
+inline constexpr Dtype bfloat16{Dtype::Val::bfloat16, sizeof(uint16_t)};
+inline constexpr Dtype complex64{Dtype::Val::complex64, sizeof(complex64_t)};
 
 Dtype promote_types(const Dtype& t1, const Dtype& t2);
 

mlx/fast.cpp

@@ -46,6 +46,141 @@ std::pair<std::vector<array>, std::vector<int>> Custom::vmap(
   return {outputs, out_axes};
 }
 
+array rms_norm(
+    const array& x,
+    const array& weight,
+    float eps,
+    StreamOrDevice s_ /* = {} */) {
+  if (x.ndim() == 0) {
+    std::ostringstream msg;
+    msg << "[rms_norm] Input must have at least 1 dimension but got input with "
+           "0 dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (weight.ndim() != 1) {
+    std::ostringstream msg;
+    msg << "[rms_norm] weight must have 1 dimension but has " << weight.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  auto out_type = result_type({x, weight});
+  if (!is_floating_point(out_type) || is_complex(out_type)) {
+    std::ostringstream msg;
+    msg << "[rms_norm] Received unsupported type " << out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+  auto fallback = [eps, out_type, s](const std::vector<array>& inputs) {
+    auto x = astype(inputs[0], float32, s);
+    x = multiply(
+        x,
+        rsqrt(
+            add(mean(square(x, s), -1, /* keepdims */ true, s),
+                array(eps, float32),
+                s),
+            s),
+        s);
+    x = astype(x, out_type, s);
+    return std::vector<array>{multiply(inputs[1], x, s)};
+  };
+  if (s.device == Device::gpu) {
+    return array(
+        x.shape(),
+        out_type,
+        std::make_unique<RMSNorm>(s, fallback, eps),
+        {astype(x, out_type, s), astype(weight, out_type, s)});
+  }
+  return fallback({x, weight})[0];
+}
+
+bool RMSNorm::is_equivalent(const Primitive& other) const {
+  const RMSNorm& a_other = static_cast<const RMSNorm&>(other);
+  return eps_ == a_other.eps_;
+}
+
+array layer_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    const std::optional<array>& bias,
+    float eps,
+    StreamOrDevice s_ /* = {} */) {
+  if (x.ndim() == 0) {
+    std::ostringstream msg;
+    msg << "[layer_norm] Input must have at least 1 dimension but got input with "
+           "0 dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (weight.has_value() && (*weight).ndim() != 1) {
+    std::ostringstream msg;
+    msg << "[layer_norm] weight must have 1 dimension but has "
+        << (*weight).ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (bias.has_value() && (*bias).ndim() != 1) {
+    std::ostringstream msg;
+    msg << "[layer_norm] bias must have 1 dimension but has " << (*bias).ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto out_type = (weight.has_value())
+      ? ((bias.has_value()) ? result_type({x, *weight, *bias})
+                            : result_type({x, *weight}))
+      : x.dtype();
+  if (!is_floating_point(out_type) || is_complex(out_type)) {
+    std::ostringstream msg;
+    msg << "[layer_norm] Received unsupported type " << out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+  bool has_weight = weight.has_value();
+  bool has_bias = bias.has_value();
+  auto fallback = [has_weight, has_bias, eps, out_type, s](
+                      const std::vector<array>& inputs) {
+    auto x = astype(inputs[0], float32, s);
+
+    // Should I not be smart here and leave the double mean to simplify()?
+    auto mu = mean(x, /* axis= */ -1, /* keepdims= */ true, s);
+    auto mu2 = square(mu, s);
+    auto x2 = mean(square(x, s), /* axis= */ -1, /* keepdims= */ true, s);
+    auto v = subtract(x2, mu2, s);
+
+    x = multiply(subtract(x, mu, s), rsqrt(add(v, array(eps, float32), s), s));
+    x = astype(x, out_type, s);
+
+    // If the LN is affine then transform x according to the weight and bias
+    if (has_weight) {
+      x = multiply(x, inputs[1], s);
+    }
+    if (has_bias) {
+      x = add(x, inputs[2], s);
+    }
+
+    return std::vector<array>{x};
+  };
+
+  auto passed_weight =
+      astype((weight.has_value()) ? *weight : array(1, out_type), out_type);
+  auto passed_bias =
+      astype((bias.has_value()) ? *bias : array(0, out_type), out_type);
+
+  if (s.device == Device::gpu) {
+    return array(
+        x.shape(),
+        out_type,
+        std::make_unique<LayerNorm>(s, fallback, eps),
+        {astype(x, out_type, s), passed_weight, passed_bias});
+  }
+  return fallback({x, passed_weight, passed_bias})[0];
+}
+
+bool LayerNorm::is_equivalent(const Primitive& other) const {
+  const LayerNorm& a_other = static_cast<const LayerNorm&>(other);
+  return eps_ == a_other.eps_;
+}
+
 array rope(
     const array& x,
     int dims,
@@ -54,10 +189,10 @@ array rope(
     float scale,
     int offset,
     StreamOrDevice s /* = {} */) {
-  if (x.ndim() != 3) {
+  if (x.ndim() < 3) {
     std::ostringstream msg;
-    msg << "[rope] Input must have 3 dimensions but got input with " << x.ndim()
-        << " dimensions.";
+    msg << "[rope] Input must have at least 3 dimensions but got input with "
+        << x.ndim() << " dimensions.";
     throw std::invalid_argument(msg.str());
   }
   if (traditional && x.shape(-1) != dims) {
@@ -67,7 +202,9 @@ array rope(
 
   auto fallback = [dims, traditional, base, scale, offset, s](
                       const std::vector<array>& inputs) {
-    auto& x = inputs[0];
+    auto& shape = inputs[0].shape();
+    int ndim = shape.size();
+    auto x = reshape(inputs[0], {-1, shape[ndim - 2], shape[ndim - 1]}, s);
     auto t = x.dtype();
     auto N = x.shape(1) + offset;
     // Compute sines and cosines
@@ -89,7 +226,7 @@ array rope(
       for (auto& o : outs) {
         o = expand_dims(o, 3, s);
       }
-      return std::vector<array>{reshape(concatenate(outs, 3, s), x.shape(), s)};
+      return std::vector<array>{reshape(concatenate(outs, 3, s), shape, s)};
     } else {
       auto out_s = x.shape();
       out_s.back() = half_dims;
@@ -103,7 +240,7 @@ array rope(
       if (dims < x.shape(-1)) {
         outs.push_back(slice(x, {0, 0, dims}, x.shape(), s));
       }
-      return std::vector<array>{concatenate(outs, 2, s)};
+      return std::vector<array>{reshape(concatenate(outs, 2, s), shape, s)};
     }
   };
   auto stream = to_stream(s);

mlx/fast.h

@@ -8,6 +8,19 @@
 
 namespace mlx::core::fast {
 
+array rms_norm(
+    const array& x,
+    const array& weight,
+    float eps,
+    StreamOrDevice s = {});
+
+array layer_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    const std::optional<array>& bias,
+    float eps,
+    StreamOrDevice s = {});
+
 array rope(
     const array& x,
     int dims,
@@ -15,7 +28,7 @@ array rope(
     float base,
     float scale,
     int offset,
-    StreamOrDevice s /* = {} */);
+    StreamOrDevice s = {});
 
 /** Computes: O = softmax(Q @ K.T) @ V **/
 array scaled_dot_product_attention(

mlx/fast_primitives.h

@@ -1,3 +1,5 @@
+// Copyright © 2024 Apple Inc.
+
 #include "mlx/primitives.h"
 
 namespace mlx::core::fast {
@@ -31,6 +33,52 @@ class Custom : public Primitive {
   std::function<std::vector<array>(std::vector<array>)> fallback_;
 };
 
+class RMSNorm : public Custom {
+ public:
+  RMSNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, fallback), eps_(eps){};
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  };
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_PRINT(RMSNorm)
+  bool is_equivalent(const Primitive& other) const override;
+
+ private:
+  std::function<std::vector<array>(std::vector<array>)> fallback_;
+  float eps_;
+};
+
+class LayerNorm : public Custom {
+ public:
+  LayerNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, fallback), eps_(eps){};
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  };
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_PRINT(LayerNorm)
+  bool is_equivalent(const Primitive& other) const override;
+
+ private:
+  std::function<std::vector<array>(std::vector<array>)> fallback_;
+  float eps_;
+};
+
 class RoPE : public Custom {
  public:
   RoPE(
@@ -49,7 +97,9 @@ class RoPE : public Custom {
         offset_(offset){};
 
   void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
-      override;
+      override {
+    throw std::runtime_error("NYI");
+  };
   void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override;
 

mlx/io/load.cpp

@@ -1,5 +1,4 @@
-// Copyright © 2023 Apple Inc.
-
+// Copyright © 2023-2024 Apple Inc.
 #include <algorithm>
 #include <cstring>
 #include <fstream>
@@ -18,7 +17,7 @@ namespace mlx::core {
 
 namespace {
 
-static constexpr uint8_t MAGIC[] = {
+constexpr uint8_t MAGIC[] = {
     0x93,
     0x4e,
     0x55,
@@ -122,8 +121,6 @@ void save(std::shared_ptr<io::Writer> out_stream, array a) {
   out_stream->write(magic_ver_len.str().c_str(), magic_ver_len.str().length());
   out_stream->write(header.str().c_str(), header.str().length());
   out_stream->write(a.data<char>(), a.nbytes());
-
-  return;
 }
 
 /** Save array to file in .npy format */

mlx/linalg.cpp

@@ -238,4 +238,27 @@ std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
       {a});
 }
 
+array inv(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() != float32) {
+    std::ostringstream msg;
+    msg << "[linalg::inv] Arrays must type float32. Received array "
+        << "with type " << a.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::inv] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(
+        "[linalg::inv] Inverses are only defined for square matrices.");
+  }
+
+  return array(
+      a.shape(), a.dtype(), std::make_unique<Inverse>(to_stream(s)), {a});
+}
+
 } // namespace mlx::core::linalg

mlx/linalg.h

@@ -64,4 +64,6 @@ std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
 
 std::vector<array> svd(const array& a, StreamOrDevice s = {});
 
+array inv(const array& a, StreamOrDevice s = {});
+
 } // namespace mlx::core::linalg

mlx/ops.cpp

@@ -445,6 +445,60 @@ array expand_dims(
   return reshape(a, out_shape, s);
 }
 
+// Slice helper
+namespace {
+
+inline auto normalize_slice(
+    const std::vector<int>& shape,
+    std::vector<int>& start,
+    std::vector<int>& stop,
+    std::vector<int>& strides) {
+  std::vector<int> out_shape(shape.size());
+  bool has_neg_strides = false;
+
+  for (int i = 0; i < shape.size(); ++i) {
+    // Following numpy docs
+    //  Negative i and j are interpreted as n + i and n + j where n is
+    //  the number of elements in the corresponding dimension. Negative
+    //  k makes stepping go towards smaller indices
+
+    auto n = shape[i];
+    auto s = start[i];
+    s = s < 0 ? s + n : s;
+    auto e = stop[i];
+    e = e < 0 ? e + n : e;
+
+    // Note: -ve strides require start >= stop
+    if (strides[i] < 0) {
+      has_neg_strides = true;
+
+      // Clamp to bounds
+      auto st = std::min(s, n - 1);
+      auto ed = std::max(-1, e);
+
+      start[i] = st;
+      stop[i] = ed > st ? st : ed;
+
+      auto str = -strides[i];
+      out_shape[i] = (start[i] - stop[i] + str - 1) / str;
+
+    } else {
+      // Clamp to bounds
+      auto st = std::max(0, std::min(s, n));
+      auto ed = std::max(0, std::min(e, n));
+
+      start[i] = st;
+      stop[i] = ed < st ? st : ed;
+
+      out_shape[i] = (stop[i] - start[i] + strides[i] - 1) / strides[i];
+    }
+  }
+
+  return std::make_pair(has_neg_strides, out_shape);
+}
+
+} // namespace
+
 array slice(
     const array& a,
     std::vector<int> start,
@@ -459,113 +513,13 @@ array slice(
     throw std::invalid_argument(msg.str());
   }
 
-  std::vector<int> negatively_strided_axes;
-  std::vector<std::vector<int>> negatively_strided_slices;
-  std::vector<int> out_shape(a.ndim());
-  for (int i = 0; i < a.ndim(); ++i) {
-    // Following numpy docs
-    //  Negative i and j are interpreted as n + i and n + j where n is
-    //  the number of elements in the corresponding dimension. Negative
-    //  k makes stepping go towards smaller indices
+  auto [has_neg_strides, out_shape] =
+      normalize_slice(a.shape(), start, stop, strides);
 
-    auto n = a.shape(i);
-    auto s = start[i];
-    s = s < 0 ? s + n : s;
-    auto e = stop[i];
-    e = e < 0 ? e + n : e;
-
-    // Note: We pass positive strides to the primitive and then flip
-    //       the axes later as needed
-    if (strides[i] < 0) {
-      negatively_strided_axes.push_back(i);
-      auto st = std::min(s, n - 1);
-      auto ed = std::max(e, -1);
-      negatively_strided_slices.push_back({st, ed, strides[i]});
-      start[i] = 0;
-      stop[i] = n;
-      strides[i] = 1;
-    } else {
-      start[i] = s;
-      stop[i] = e < s ? s : e;
-    }
-
-    // Clamp to bounds
-    start[i] = std::max(0, std::min(start[i], n));
-    stop[i] = std::max(0, std::min(stop[i], n));
-
-    out_shape[i] = (stop[i] - start[i] + strides[i] - 1) / strides[i];
-  }
-
-  // If strides are negative, slice and then make a copy with axes flipped
-  if (negatively_strided_axes.size() > 0) {
-    // First, take the slice of the positively strided axes
-    auto out = array(
-        out_shape,
-        a.dtype(),
-        std::make_unique<Slice>(
-            to_stream(s),
-            std::move(start),
-            std::move(stop),
-            std::move(strides)),
-        {a});
-
-    std::vector<array> indices;
-    std::vector<int> slice_sizes = out.shape();
-    std::vector<int> t_axes(out.ndim(), -1);
-    std::vector<int> out_reshape(out.ndim(), -1);
-
-    int n_axes = negatively_strided_axes.size();
-    for (int i = 0; i < n_axes; i++) {
-      // Get axis and corresponding slice
-      auto ax = negatively_strided_axes[i];
-      auto sl = negatively_strided_slices[i];
-
-      // Get indices for the slice
-      auto ax_idx = arange(sl[0], sl[1], sl[2], s);
-
-      // Reshape indices for broadcast as needed
-      std::vector<int> ax_idx_shape(n_axes, 1);
-      ax_idx_shape[i] = ax_idx.size();
-      ax_idx = reshape(ax_idx, ax_idx_shape, s);
-
-      // Add indices to list
-      indices.push_back(ax_idx);
-
-      // Set slice size for axis
-      slice_sizes[ax] = 1;
-
-      // Gather moves the axis up, remainder needs to be squeezed
-      out_reshape[i] = indices[i].size();
-
-      // Gather moves the axis up, needs to be transposed
-      t_axes[ax] = i;
-    }
-
-    // Prepare out_reshape to squeeze gathered dims
-    // Prepare to transpose dims as needed
-    int j = n_axes;
-    for (int i = 0; j < out.ndim() && i < out.ndim(); i++) {
-      if (t_axes[i] < 0) {
-        t_axes[i] = j;
-        out_reshape[j] = out_shape[i];
-        j++;
-      }
-    }
-
-    // Gather
-    out = gather(out, indices, negatively_strided_axes, slice_sizes, s);
-
-    // Squeeze dims
-    out = reshape(out, out_reshape, s);
-
-    // Transpose dims
-    out = transpose(out, t_axes, s);
-
-    return out;
-  }
-  if (out_shape == a.shape()) {
+  if (!has_neg_strides && out_shape == a.shape()) {
     return a;
   }
+
   return array(
       out_shape,
       a.dtype(),
@@ -582,6 +536,55 @@ array slice(
   return slice(a, start, stop, std::vector<int>(a.ndim(), 1), to_stream(s));
 }
 
+/** Update a slice from the source array */
+array slice_update(
+    const array& src,
+    const array& update,
+    std::vector<int> start,
+    std::vector<int> stop,
+    std::vector<int> strides,
+    StreamOrDevice s /* = {} */) {
+  // Check dimensions
+  if (start.size() != src.ndim() || stop.size() != src.ndim() ||
+      strides.size() != src.ndim()) {
+    std::ostringstream msg;
+    msg << "[slice] Invalid number of indices or strides for "
+        << "array with dimension " << src.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Process slice dimensions
+  auto [has_neg_strides, upd_shape] =
+      normalize_slice(src.shape(), start, stop, strides);
+
+  // Broadcast update shape to slice shape
+  auto update_broadcasted = broadcast_to(update, upd_shape, s);
+
+  // If the entire src is the slice, just return the update
+  if (!has_neg_strides && upd_shape == src.shape()) {
+    return astype(update_broadcasted, src.dtype(), s);
+  }
+
+  return array(
+      src.shape(),
+      src.dtype(),
+      std::make_unique<SliceUpdate>(
+          to_stream(s), std::move(start), std::move(stop), std::move(strides)),
+      {src, update_broadcasted});
+}
+
+/** Update a slice from the source array with stride 1 in each dimension */
+array slice_update(
+    const array& src,
+    const array& update,
+    std::vector<int> start,
+    std::vector<int> stop,
+    StreamOrDevice s /* = {} */) {
+  auto strides = std::vector<int>(src.ndim(), 1);
+  return slice_update(
+      src, update, std::move(start), std::move(stop), std::move(strides), s);
+}
+
 std::vector<array> split(
     const array& a,
     const std::vector<int>& indices,

mlx/ops.h

@@ -177,6 +177,23 @@ array slice(
     const std::vector<int>& stop,
     StreamOrDevice s = {});
 
+/** Update a slice from the source array */
+array slice_update(
+    const array& src,
+    const array& update,
+    std::vector<int> start,
+    std::vector<int> stop,
+    std::vector<int> strides,
+    StreamOrDevice s = {});
+
+/** Update a slice from the source array with stride 1 in each dimension */
+array slice_update(
+    const array& src,
+    const array& update,
+    std::vector<int> start,
+    std::vector<int> stop,
+    StreamOrDevice s = {});
+
 /** Split an array into sub-arrays along a given axis. */
 std::vector<array>
 split(const array& a, int num_splits, int axis, StreamOrDevice s = {});

mlx/primitives.cpp

@@ -110,7 +110,11 @@ std::vector<array> Primitive::jvp(
     const std::vector<array>&,
     const std::vector<array>&,
     const std::vector<int>&) {
-  throw std::invalid_argument("Primitive's jvp not implemented.");
+  std::ostringstream msg;
+  msg << "[Primitive::jvp] Not implemented for ";
+  print(msg);
+  msg << ".";
+  throw std::invalid_argument(msg.str());
 };
 
 std::vector<array> Primitive::vjp(
@@ -118,13 +122,21 @@ std::vector<array> Primitive::vjp(
     const std::vector<array>&,
     const std::vector<int>&,
     const std::vector<array>&) {
-  throw std::invalid_argument("Primitive's vjp not implemented.");
+  std::ostringstream msg;
+  msg << "[Primitive::vip] Not implemented for ";
+  print(msg);
+  msg << ".";
+  throw std::invalid_argument(msg.str());
 };
 
 std::pair<std::vector<array>, std::vector<int>> Primitive::vmap(
     const std::vector<array>&,
     const std::vector<int>&) {
-  throw std::invalid_argument("Primitive's vmap not implemented.");
+  std::ostringstream msg;
+  msg << "[Primitive::vmap] Not implemented for ";
+  print(msg);
+  msg << ".";
+  throw std::invalid_argument(msg.str());
 };
 
 std::vector<std::vector<int>> Primitive::output_shapes(
@@ -235,6 +247,18 @@ bool AddMM::is_equivalent(const Primitive& other) const {
   return (alpha_ == a_other.alpha_ && beta_ == a_other.beta_);
 }
 
+std::pair<std::vector<array>, std::vector<int>> AddMM::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto maybe_move_ax = [this](auto& arr, auto ax) {
+    return ax > 0 ? moveaxis(arr, ax, 0, stream()) : arr;
+  };
+  auto a = maybe_move_ax(inputs[0], axes[0]);
+  auto b = maybe_move_ax(inputs[1], axes[1]);
+  auto c = maybe_move_ax(inputs[2], axes[2]);
+  return {{addmm(c, a, b, alpha_, beta_, stream())}, {0}};
+}
+
 bool Arange::is_equivalent(const Primitive& other) const {
   const Arange& a_other = static_cast<const Arange&>(other);
   return (
@@ -1103,7 +1127,7 @@ std::pair<std::vector<array>, std::vector<int>> Equal::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
   auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
-  return {{equal(a, b, stream())}, axes};
+  return {{equal(a, b, stream())}, {to_ax}};
 }
 
 std::vector<array> Equal::vjp(
@@ -1444,7 +1468,7 @@ std::pair<std::vector<array>, std::vector<int>> Greater::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
   auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
-  return {{greater(a, b, stream())}, axes};
+  return {{greater(a, b, stream())}, {to_ax}};
 }
 
 std::vector<array> Greater::vjp(
@@ -1471,7 +1495,7 @@ std::pair<std::vector<array>, std::vector<int>> GreaterEqual::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
   auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
-  return {{greater_equal(a, b, stream())}, axes};
+  return {{greater_equal(a, b, stream())}, {to_ax}};
 }
 
 std::vector<array> GreaterEqual::vjp(
@@ -1498,7 +1522,7 @@ std::pair<std::vector<array>, std::vector<int>> Less::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
   auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
-  return {{less(a, b, stream())}, axes};
+  return {{less(a, b, stream())}, {to_ax}};
 }
 
 std::vector<array> Less::vjp(
@@ -1525,7 +1549,7 @@ std::pair<std::vector<array>, std::vector<int>> LessEqual::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
   auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
-  return {{less_equal(a, b, stream())}, axes};
+  return {{less_equal(a, b, stream())}, {to_ax}};
 }
 
 std::vector<array> LessEqual::vjp(
@@ -1772,6 +1796,17 @@ std::vector<array> Matmul::vjp(
   return vjps;
 }
 
+std::pair<std::vector<array>, std::vector<int>> Matmul::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto maybe_move_ax = [this](auto& arr, auto ax) {
+    return ax > 0 ? moveaxis(arr, ax, 0, stream()) : arr;
+  };
+  auto a = maybe_move_ax(inputs[0], axes[0]);
+  auto b = maybe_move_ax(inputs[1], axes[1]);
+  return {{matmul(a, b, stream())}, {0}};
+}
+
 std::vector<array> Maximum::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
@@ -2814,6 +2849,114 @@ bool Slice::is_equivalent(const Primitive& other) const {
       end_indices_ == s_other.end_indices_ && strides_ == s_other.strides_);
 }
 
+std::pair<std::vector<array>, std::vector<int>> SliceUpdate::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+
+  auto start = start_indices_;
+  auto stop = end_indices_;
+  auto strides = strides_;
+
+  auto src = inputs[0];
+  auto upd = inputs[1];
+
+  auto src_ax = axes[0];
+  auto upd_ax = axes[1];
+
+  // No vmapping needed
+  if (src_ax == -1 && upd_ax == -1) {
+    return {{slice_update(src, upd, start, stop, strides, stream())}, {-1}};
+  }
+
+  // Broadcast src
+  if (src_ax == -1) {
+    src = expand_dims(src, upd_ax, stream());
+    auto shape = src.shape();
+    shape[upd_ax] = upd.shape(upd_ax);
+    src = broadcast_to(src, shape, stream());
+    src_ax = upd_ax;
+  }
+
+  // Broadcast upd
+  if (upd_ax == -1) {
+    upd = expand_dims(upd, src_ax, stream());
+    upd_ax = src_ax;
+  }
+
+  if (src_ax != upd_ax) {
+    upd = moveaxis(upd, upd_ax, src_ax, stream());
+  }
+
+  start.insert(start.begin() + src_ax, 0);
+  stop.insert(stop.begin() + src_ax, src.shape(src_ax));
+  strides.insert(strides.begin() + src_ax, 1);
+
+  return {{slice_update(src, upd, start, stop, strides, stream())}, {src_ax}};
+}
+
+std::vector<array> SliceUpdate::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  // Check inputs
+  assert(primals.size() == 2);
+
+  auto& cotan = cotangents[0];
+  auto& src = primals[0];
+  auto& upd = primals[1];
+
+  std::vector<array> vjps;
+
+  for (int num : argnums) {
+    // Vjp for source
+    if (num == 0) {
+      auto grad = slice_update(
+          cotan,
+          zeros_like(upd, stream()),
+          start_indices_,
+          end_indices_,
+          strides_,
+          stream());
+
+      vjps.push_back(grad);
+    }
+    // Vjp fpr updates
+    else {
+      auto grad =
+          slice(cotan, start_indices_, end_indices_, strides_, stream());
+
+      vjps.push_back(grad);
+    }
+  }
+
+  return vjps;
+}
+
+std::vector<array> SliceUpdate::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  // Check inputs
+  assert(primals.size() == 2);
+  return {slice_update(
+      tangents[0],
+      tangents[1],
+      start_indices_,
+      end_indices_,
+      strides_,
+      stream())};
+}
+
+bool SliceUpdate::is_equivalent(const Primitive& other) const {
+  const SliceUpdate& s_other = static_cast<const SliceUpdate&>(other);
+  return (
+      start_indices_ == s_other.start_indices_ &&
+      end_indices_ == s_other.end_indices_ && strides_ == s_other.strides_);
+}
+
 std::pair<std::vector<array>, std::vector<int>> Softmax::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {

mlx/primitives.h

@@ -200,6 +200,7 @@ class AddMM : public UnaryPrimitive {
       const std::vector<int>& argnums,
       const std::vector<array>& outputs) override;
 
+  DEFINE_VMAP()
   DEFINE_PRINT(AddMM)
 
   bool is_equivalent(const Primitive& other) const override;
@@ -1140,6 +1141,7 @@ class Matmul : public UnaryPrimitive {
       const std::vector<int>& argnums,
       const std::vector<array>& outputs) override;
 
+  DEFINE_VMAP()
   DEFINE_PRINT(Matmul)
   DEFINE_DEFAULT_IS_EQUIVALENT()
 };
@@ -1658,6 +1660,44 @@ class Slice : public UnaryPrimitive {
   std::vector<int> strides_;
 
   void eval(const std::vector<array>& inputs, array& out);
+
+  std::tuple<bool, int64_t, std::vector<int64_t>> prepare_slice(
+      const array& in);
+  void shared_buffer_slice(
+      const array& in,
+      const std::vector<size_t>& out_strides,
+      size_t data_offset,
+      array& out);
+};
+
+class SliceUpdate : public UnaryPrimitive {
+ public:
+  explicit SliceUpdate(
+      Stream stream,
+      const std::vector<int>& start_indices,
+      const std::vector<int>& end_indices,
+      const std::vector<int>& strides)
+      : UnaryPrimitive(stream),
+        start_indices_(start_indices),
+        end_indices_(end_indices),
+        strides_(strides){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_PRINT(SliceUpdate)
+  bool is_equivalent(const Primitive& other) const override;
+
+ private:
+  std::vector<int> start_indices_;
+  std::vector<int> end_indices_;
+  std::vector<int> strides_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+
+  std::tuple<int64_t, std::vector<int64_t>> prepare_slice(const array& in);
 };
 
 class Softmax : public UnaryPrimitive {
@@ -1889,10 +1929,26 @@ class SVD : public Primitive {
   void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override;
 
+  DEFINE_VMAP()
   DEFINE_PRINT(SVD)
 
  private:
   void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
 };
 
+/* Matrix inversion primitive. */
+class Inverse : public UnaryPrimitive {
+ public:
+  explicit Inverse(Stream stream) : UnaryPrimitive(stream){};
+
+  void eval_cpu(const std::vector<array>& inputs, array& output) override;
+  void eval_gpu(const std::vector<array>& inputs, array& output) override;
+
+  DEFINE_VMAP()
+  DEFINE_PRINT(Inverse)
+
+ private:
+  void eval(const std::vector<array>& inputs, array& output);
+};
+
 } // namespace mlx::core

mlx/transforms.cpp

@@ -37,7 +37,7 @@ class Synchronizer : public Primitive {
 // are currently under a function transformation.
 int detail::InTracing::tracing_counter{0};
 
-void eval(const std::vector<array>& outputs) {
+void eval(std::vector<array> outputs) {
   std::function<void(const array&)> recurse;
   std::queue<array> tape;
   std::unordered_set<std::uintptr_t> cache;
@@ -52,8 +52,8 @@ void eval(const std::vector<array>& outputs) {
     }
   }
 
-  auto synchronizer =
-      array({}, bool_, std::make_unique<Synchronizer>(stream), outputs);
+  auto synchronizer = array(
+      {}, bool_, std::make_unique<Synchronizer>(stream), std::move(outputs));
 
   size_t depth_counter = 0;
   recurse = [&](const array& a) {
@@ -118,7 +118,7 @@ void eval(const std::vector<array>& outputs) {
         arr_deps.push_back(it->second);
       }
     }
-    std::shared_ptr<std::promise<void>> p{nullptr};
+    std::shared_ptr<std::promise<void>> p;
     if (auto it = deps.find(arr.primitive_id()); it != deps.end()) {
       p = std::make_unique<std::promise<void>>();
       ps.push_back(p);
@@ -655,6 +655,7 @@ std::vector<array> vmap_replace(
     }
 
     auto [v_outputs, v_out_axes] = a.primitive().vmap(v_inputs, v_axes);
+
     // For each primitive's outputs add its id, the vout id and the vax
     auto outputs = a.outputs();
     for (int i = 0; i < v_outputs.size(); ++i) {

mlx/transforms.h

@@ -6,10 +6,10 @@
 
 namespace mlx::core {
 
-void eval(const std::vector<array>& outputs);
+void eval(std::vector<array> outputs);
 
 template <typename... Arrays>
-void eval(Arrays... outputs) {
+void eval(Arrays&&... outputs) {
   eval(std::vector<array>{std::forward<Arrays>(outputs)...});
 }
 

mlx/types/complex.h

@@ -7,9 +7,28 @@
 namespace mlx::core {
 
 struct complex64_t;
+struct complex128_t;
 
 template <typename T>
-static constexpr bool can_convert_to_complex64 =
+inline constexpr bool can_convert_to_complex128 =
+    !std::is_same_v<T, complex128_t> && std::is_convertible_v<T, double>;
+
+struct complex128_t : public std::complex<double> {
+  complex128_t(double v, double u) : std::complex<double>(v, u){};
+  complex128_t(std::complex<double> v) : std::complex<double>(v){};
+
+  template <
+      typename T,
+      typename = typename std::enable_if<can_convert_to_complex128<T>>::type>
+  complex128_t(T x) : std::complex<double>(x){};
+
+  operator float() const {
+    return real();
+  };
+};
+
+template <typename T>
+inline constexpr bool can_convert_to_complex64 =
     !std::is_same_v<T, complex64_t> && std::is_convertible_v<T, float>;
 
 struct complex64_t : public std::complex<float> {

mlx/utils.cpp

@@ -329,4 +329,13 @@ std::ostream& operator<<(std::ostream& os, const std::vector<size_t>& v) {
   return os;
 }
 
+std::ostream& operator<<(std::ostream& os, const std::vector<int64_t>& v) {
+  os << "(";
+  for (int i = 0; i < v.size(); ++i) {
+    os << v[i] << ((i == v.size() - 1) ? "" : ",");
+  }
+  os << ")";
+  return os;
+}
+
 } // namespace mlx::core

mlx/utils.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -76,6 +76,7 @@ std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k);
 std::ostream& operator<<(std::ostream& os, array a);
 std::ostream& operator<<(std::ostream& os, const std::vector<int>& v);
 std::ostream& operator<<(std::ostream& os, const std::vector<size_t>& v);
+std::ostream& operator<<(std::ostream& os, const std::vector<int64_t>& v);
 inline std::ostream& operator<<(std::ostream& os, const complex64_t& v) {
   return os << v.real() << (v.imag() >= 0 ? "+" : "") << v.imag() << "j";
 }

pyproject.toml

@@ -1,3 +1,7 @@
 [build-system]
-requires = ["setuptools>=42", "pybind11>=2.10", "cmake>=3.24"]
+requires = [
+  "setuptools>=42",
+  "nanobind@git+https://github.com/wjakob/nanobind.git#egg=4148debcf91f5ccab0c3b8d67b5c3cabd61f407f",
+  "cmake>=3.24",
+]
 build-backend = "setuptools.build_meta"

python/mlx/nn/layers/normalization.py

@@ -85,13 +85,19 @@ class LayerNorm(Module):
         eps (float): A small additive constant for numerical stability
         affine (bool): If True learn an affine transform to apply after the
             normalization
+        bias (bool): If True include a translation to the affine
+            transformation. If set to False the transformation is not really affine
+            just scaling.
     """
 
-    def __init__(self, dims: int, eps: float = 1e-5, affine: bool = True):
+    def __init__(
+        self, dims: int, eps: float = 1e-5, affine: bool = True, bias: bool = True
+    ):
         super().__init__()
         if affine:
-            self.bias = mx.zeros((dims,))
             self.weight = mx.ones((dims,))
+            if bias:
+                self.bias = mx.zeros((dims,))
         self.eps = eps
         self.dims = dims
 
@@ -99,10 +105,9 @@ class LayerNorm(Module):
         return f"{self.dims}, eps={self.eps}, affine={'weight' in self}"
 
     def __call__(self, x):
-        means = mx.mean(x, axis=-1, keepdims=True)
-        var = mx.var(x, axis=-1, keepdims=True)
-        x = (x - means) * mx.rsqrt(var + self.eps)
-        return (self.weight * x + self.bias) if "weight" in self else x
+        weight = self.weight if "weight" in self else None
+        bias = self.bias if "bias" in self else None
+        return mx.fast.layer_norm(x, weight, bias, self.eps)
 
 
 class RMSNorm(Module):
@@ -117,6 +122,8 @@ class RMSNorm(Module):
     where :math:`\gamma` is a learned per feature dimension parameter initialized at
     1.
 
+    Note the accumulation for the mean is done in 32-bit precision.
+
     [1]: https://arxiv.org/abs/1910.07467
 
     Args:
@@ -133,18 +140,7 @@ class RMSNorm(Module):
         return f"{self.weight.shape[0]}, eps={self.eps}"
 
     def __call__(self, x):
-        # S is 1/sqrt(N) where N is the size of the features of x and is used
-        # to compute a numerically more stable RMS of x by multiplying with S
-        # first and summing.
-        #
-        # This way we prefer underflow over overflow which is controlled with
-        # the parameter epsilon anyway.
-        S = 1 / x.shape[-1] ** 0.5
-
-        n = (x * S).square().sum(axis=-1, keepdims=True)
-        n = mx.rsqrt(n + self.eps)
-
-        return self.weight * x * n
+        return mx.fast.rms_norm(x, self.weight, self.eps)
 
 
 class GroupNorm(Module):

python/mlx/nn/layers/pooling.py

@@ -20,6 +20,22 @@ def _value_or_list(x, n, msg):
     return [x] * n
 
 
+def _non_overlapping_sliding_windows(x, shape, window_shape):
+    # Compute the intermediate shape
+    new_shape = [shape[0]]
+    for s, w in zip(shape[1:], window_shape):
+        new_shape.append(s // w)
+        new_shape.append(w)
+    new_shape.append(shape[-1])
+
+    last_axis = len(new_shape) - 1
+    axis_order = [0, *range(1, last_axis, 2), *range(2, last_axis, 2), last_axis]
+
+    x = x.reshape(new_shape)
+    x = x.transpose(axis_order)
+    return x
+
+
 def _sliding_windows(x, window_shape, window_strides):
     if x.ndim < 3:
         raise ValueError(
@@ -37,6 +53,12 @@ def _sliding_windows(x, window_shape, window_strides):
         )
 
     shape = x.shape
+    if all(
+        window == stride and size % window == 0
+        for size, window, stride in zip(spatial_dims, window_shape, window_strides)
+    ):
+        return _non_overlapping_sliding_windows(x, shape, window_shape)
+
     strides = list(reversed(list(accumulate(reversed(shape + (1,)), operator.mul))))[1:]
 
     # Compute the output shape

python/mlx/nn/layers/positional_encoding.py

@@ -44,9 +44,7 @@ class RoPE(Module):
         return f"{self.dims}, traditional={self.traditional}"
 
     def __call__(self, x, offset: int = 0):
-        shape = x.shape
-        x = mx.reshape(x, (-1, shape[-2], shape[-1]))
-        x = mx.fast.rope(
+        return mx.fast.rope(
             x,
             self.dims,
             traditional=self.traditional,
@@ -54,7 +52,6 @@ class RoPE(Module):
             scale=self.scale,
             offset=offset,
         )
-        return mx.reshape(x, shape)
 
 
 class SinusoidalPositionalEncoding(Module):

python/mlx/optimizers/schedulers.py

@@ -58,13 +58,14 @@ def step_decay(init: float, decay_rate: float, step_size: int) -> Callable:
     return schedule
 
 
-def cosine_decay(init: float, decay_steps: int) -> Callable:
+def cosine_decay(init: float, decay_steps: int, minimum: float = 0.0) -> Callable:
     r"""Make a cosine decay scheduler.
 
     Args:
         init (float): Initial value.
         decay_steps (int): Number of steps to decay over. The decayed
             value is constant for steps beyond ``decay_steps``.
+        minimum (float, optional): Minimal value to decay to. Default: ``0``.
 
     Example:
 
@@ -82,7 +83,7 @@ def cosine_decay(init: float, decay_steps: int) -> Callable:
     def scheduler(step):
         s = mx.minimum(step, decay_steps)
         decay = 0.5 * (1.0 + mx.cos((math.pi / decay_steps) * s))
-        return init * decay
+        return mx.maximum(init * decay, minimum)
 
     return scheduler
 

python/src/CMakeLists.txt

@@ -1,7 +1,10 @@
-pybind11_add_module(
+nanobind_add_module(
   core
+  NB_STATIC STABLE_ABI LTO NOMINSIZE
+  NB_DOMAIN mlx
   ${CMAKE_CURRENT_SOURCE_DIR}/mlx.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/convert.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
@@ -15,7 +18,6 @@ pybind11_add_module(
   ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/constants.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/trees.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
 )
 
 if (NOT MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY)

python/src/buffer.h

@@ -0,0 +1,122 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+#include <optional>
+
+#include <nanobind/nanobind.h>
+
+#include "mlx/array.h"
+#include "mlx/utils.h"
+
+// Only defined in >= Python 3.9
+// https://github.com/python/cpython/blob/f6cdc6b4a191b75027de342aa8b5d344fb31313e/Include/typeslots.h#L2-L3
+#ifndef Py_bf_getbuffer
+#define Py_bf_getbuffer 1
+#define Py_bf_releasebuffer 2
+#endif
+
+namespace nb = nanobind;
+using namespace mlx::core;
+
+std::string buffer_format(const array& a) {
+  // https://docs.python.org/3.10/library/struct.html#format-characters
+  switch (a.dtype()) {
+    case bool_:
+      return "?";
+    case uint8:
+      return "B";
+    case uint16:
+      return "H";
+    case uint32:
+      return "I";
+    case uint64:
+      return "Q";
+    case int8:
+      return "b";
+    case int16:
+      return "h";
+    case int32:
+      return "i";
+    case int64:
+      return "q";
+    case float16:
+      return "e";
+    case float32:
+      return "f";
+    case bfloat16:
+      return "B";
+    case complex64:
+      return "Zf\0";
+    default: {
+      std::ostringstream os;
+      os << "bad dtype: " << a.dtype();
+      throw std::runtime_error(os.str());
+    }
+  }
+}
+
+struct buffer_info {
+  std::string format;
+  std::vector<ssize_t> shape;
+  std::vector<ssize_t> strides;
+
+  buffer_info(
+      const std::string& format,
+      std::vector<ssize_t> shape_in,
+      std::vector<ssize_t> strides_in)
+      : format(format),
+        shape(std::move(shape_in)),
+        strides(std::move(strides_in)) {}
+
+  buffer_info(const buffer_info&) = delete;
+  buffer_info& operator=(const buffer_info&) = delete;
+
+  buffer_info(buffer_info&& other) noexcept {
+    (*this) = std::move(other);
+  }
+
+  buffer_info& operator=(buffer_info&& rhs) noexcept {
+    format = std::move(rhs.format);
+    shape = std::move(rhs.shape);
+    strides = std::move(rhs.strides);
+    return *this;
+  }
+};
+
+extern "C" inline int getbuffer(PyObject* obj, Py_buffer* view, int flags) {
+  std::memset(view, 0, sizeof(Py_buffer));
+  auto a = nb::cast<array>(nb::handle(obj));
+
+  if (!a.is_evaled()) {
+    nb::gil_scoped_release nogil;
+    a.eval();
+  }
+
+  std::vector<ssize_t> shape(a.shape().begin(), a.shape().end());
+  std::vector<ssize_t> strides(a.strides().begin(), a.strides().end());
+  for (auto& s : strides) {
+    s *= a.itemsize();
+  }
+  buffer_info* info =
+      new buffer_info(buffer_format(a), std::move(shape), std::move(strides));
+
+  view->obj = obj;
+  view->ndim = a.ndim();
+  view->internal = info;
+  view->buf = a.data<void>();
+  view->itemsize = a.itemsize();
+  view->len = a.size();
+  view->readonly = false;
+  if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT) {
+    view->format = const_cast<char*>(info->format.c_str());
+  }
+  if ((flags & PyBUF_STRIDES) == PyBUF_STRIDES) {
+    view->strides = info->strides.data();
+    view->shape = info->shape.data();
+  }
+  Py_INCREF(view->obj);
+  return 0;
+}
+
+extern "C" inline void releasebuffer(PyObject*, Py_buffer* view) {
+  delete (buffer_info*)view->internal;
+}

python/src/constants.cpp

@@ -1,11 +1,11 @@
-// init_constants.cpp
+// Copyright © 2023-2024 Apple Inc.
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
 #include <limits>
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
-void init_constants(py::module_& m) {
+void init_constants(nb::module_& m) {
   m.attr("Inf") = std::numeric_limits<double>::infinity();
   m.attr("Infinity") = std::numeric_limits<double>::infinity();
   m.attr("NAN") = NAN;
@@ -19,6 +19,6 @@ void init_constants(py::module_& m) {
   m.attr("inf") = std::numeric_limits<double>::infinity();
   m.attr("infty") = std::numeric_limits<double>::infinity();
   m.attr("nan") = NAN;
-  m.attr("newaxis") = pybind11::none();
+  m.attr("newaxis") = nb::none();
   m.attr("pi") = 3.1415926535897932384626433;
-}
+}

python/src/convert.cpp

@@ -0,0 +1,155 @@
+// Copyright © 2024 Apple Inc.
+
+#include <nanobind/stl/complex.h>
+
+#include "python/src/convert.h"
+
+namespace nanobind {
+template <>
+struct ndarray_traits<float16_t> {
+  static constexpr bool is_complex = false;
+  static constexpr bool is_float = true;
+  static constexpr bool is_bool = false;
+  static constexpr bool is_int = false;
+  static constexpr bool is_signed = true;
+};
+
+template <>
+struct ndarray_traits<bfloat16_t> {
+  static constexpr bool is_complex = false;
+  static constexpr bool is_float = true;
+  static constexpr bool is_bool = false;
+  static constexpr bool is_int = false;
+  static constexpr bool is_signed = true;
+};
+
+static constexpr dlpack::dtype bfloat16{4, 16, 1};
+}; // namespace nanobind
+
+template <typename T>
+array nd_array_to_mlx_contiguous(
+    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    const std::vector<int>& shape,
+    Dtype dtype) {
+  // Make a copy of the numpy buffer
+  // Get buffer ptr pass to array constructor
+  auto data_ptr = nd_array.data();
+  return array(static_cast<const T*>(data_ptr), shape, dtype);
+}
+
+array nd_array_to_mlx(
+    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    std::optional<Dtype> dtype) {
+  // Compute the shape and size
+  std::vector<int> shape;
+  for (int i = 0; i < nd_array.ndim(); i++) {
+    shape.push_back(nd_array.shape(i));
+  }
+  auto type = nd_array.dtype();
+
+  // Copy data and make array
+  if (type == nb::dtype<bool>()) {
+    return nd_array_to_mlx_contiguous<bool>(
+        nd_array, shape, dtype.value_or(bool_));
+  } else if (type == nb::dtype<uint8_t>()) {
+    return nd_array_to_mlx_contiguous<uint8_t>(
+        nd_array, shape, dtype.value_or(uint8));
+  } else if (type == nb::dtype<uint16_t>()) {
+    return nd_array_to_mlx_contiguous<uint16_t>(
+        nd_array, shape, dtype.value_or(uint16));
+  } else if (type == nb::dtype<uint32_t>()) {
+    return nd_array_to_mlx_contiguous<uint32_t>(
+        nd_array, shape, dtype.value_or(uint32));
+  } else if (type == nb::dtype<uint64_t>()) {
+    return nd_array_to_mlx_contiguous<uint64_t>(
+        nd_array, shape, dtype.value_or(uint64));
+  } else if (type == nb::dtype<int8_t>()) {
+    return nd_array_to_mlx_contiguous<int8_t>(
+        nd_array, shape, dtype.value_or(int8));
+  } else if (type == nb::dtype<int16_t>()) {
+    return nd_array_to_mlx_contiguous<int16_t>(
+        nd_array, shape, dtype.value_or(int16));
+  } else if (type == nb::dtype<int32_t>()) {
+    return nd_array_to_mlx_contiguous<int32_t>(
+        nd_array, shape, dtype.value_or(int32));
+  } else if (type == nb::dtype<int64_t>()) {
+    return nd_array_to_mlx_contiguous<int64_t>(
+        nd_array, shape, dtype.value_or(int64));
+  } else if (type == nb::dtype<float16_t>()) {
+    return nd_array_to_mlx_contiguous<float16_t>(
+        nd_array, shape, dtype.value_or(float16));
+  } else if (type == nb::dtype<bfloat16_t>()) {
+    return nd_array_to_mlx_contiguous<bfloat16_t>(
+        nd_array, shape, dtype.value_or(bfloat16));
+  } else if (type == nb::dtype<float>()) {
+    return nd_array_to_mlx_contiguous<float>(
+        nd_array, shape, dtype.value_or(float32));
+  } else if (type == nb::dtype<double>()) {
+    return nd_array_to_mlx_contiguous<double>(
+        nd_array, shape, dtype.value_or(float32));
+  } else if (type == nb::dtype<std::complex<float>>()) {
+    return nd_array_to_mlx_contiguous<complex64_t>(
+        nd_array, shape, dtype.value_or(complex64));
+  } else if (type == nb::dtype<std::complex<double>>()) {
+    return nd_array_to_mlx_contiguous<complex128_t>(
+        nd_array, shape, dtype.value_or(complex64));
+  } else {
+    throw std::invalid_argument("Cannot convert numpy array to mlx array.");
+  }
+}
+
+template <typename Lib, typename T>
+nb::ndarray<Lib> mlx_to_nd_array(
+    array a,
+    std::optional<nb::dlpack::dtype> t = {}) {
+  // Eval if not already evaled
+  if (!a.is_evaled()) {
+    nb::gil_scoped_release nogil;
+    a.eval();
+  }
+  std::vector<size_t> shape(a.shape().begin(), a.shape().end());
+  std::vector<int64_t> strides(a.strides().begin(), a.strides().end());
+  return nb::ndarray<Lib>(
+      a.data<T>(),
+      a.ndim(),
+      shape.data(),
+      nb::handle(),
+      strides.data(),
+      t.value_or(nb::dtype<T>()));
+}
+
+template <typename Lib>
+nb::ndarray<Lib> mlx_to_nd_array(const array& a) {
+  switch (a.dtype()) {
+    case bool_:
+      return mlx_to_nd_array<Lib, bool>(a);
+    case uint8:
+      return mlx_to_nd_array<Lib, uint8_t>(a);
+    case uint16:
+      return mlx_to_nd_array<Lib, uint16_t>(a);
+    case uint32:
+      return mlx_to_nd_array<Lib, uint32_t>(a);
+    case uint64:
+      return mlx_to_nd_array<Lib, uint64_t>(a);
+    case int8:
+      return mlx_to_nd_array<Lib, int8_t>(a);
+    case int16:
+      return mlx_to_nd_array<Lib, int16_t>(a);
+    case int32:
+      return mlx_to_nd_array<Lib, int32_t>(a);
+    case int64:
+      return mlx_to_nd_array<Lib, int64_t>(a);
+    case float16:
+      return mlx_to_nd_array<Lib, float16_t>(a);
+    case bfloat16:
+      return mlx_to_nd_array<Lib, bfloat16_t>(a, nb::bfloat16);
+    case float32:
+      return mlx_to_nd_array<Lib, float>(a);
+    case complex64:
+      return mlx_to_nd_array<Lib, std::complex<float>>(a);
+  }
+}
+
+nb::ndarray<nb::numpy> mlx_to_np_array(const array& a) {
+  return mlx_to_nd_array<nb::numpy>(a);
+}

python/src/convert.h

@@ -0,0 +1,16 @@
+// Copyright © 2024 Apple Inc.
+
+#include <optional>
+
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+
+#include "mlx/array.h"
+
+namespace nb = nanobind;
+using namespace mlx::core;
+
+array nd_array_to_mlx(
+    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
+    std::optional<Dtype> dtype);
+nb::ndarray<nb::numpy> mlx_to_np_array(const array& a);

python/src/device.cpp

@@ -1,32 +1,34 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <sstream>
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
 
 #include "mlx/device.h"
 #include "mlx/utils.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 
-void init_device(py::module_& m) {
-  auto device_class = py::class_<Device>(
+void init_device(nb::module_& m) {
+  auto device_class = nb::class_<Device>(
       m, "Device", R"pbdoc(A device to run operations on.)pbdoc");
-  py::enum_<Device::DeviceType>(m, "DeviceType")
+  nb::enum_<Device::DeviceType>(m, "DeviceType")
       .value("cpu", Device::DeviceType::cpu)
       .value("gpu", Device::DeviceType::gpu)
       .export_values()
-      .def(
-          "__eq__",
-          [](const Device::DeviceType& d1, const Device& d2) {
-            return d1 == d2;
-          },
-          py::prepend());
+      .def("__eq__", [](const Device::DeviceType& d, const nb::object& other) {
+        if (!nb::isinstance<Device>(other) &&
+            !nb::isinstance<Device::DeviceType>(other)) {
+          return false;
+        }
+        return d == nb::cast<Device>(other);
+      });
 
-  device_class.def(py::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
-      .def_readonly("type", &Device::type)
+  device_class.def(nb::init<Device::DeviceType, int>(), "type"_a, "index"_a = 0)
+      .def_ro("type", &Device::type)
       .def(
           "__repr__",
           [](const Device& d) {
@@ -34,11 +36,15 @@ void init_device(py::module_& m) {
             os << d;
             return os.str();
           })
-      .def("__eq__", [](const Device& d1, const Device& d2) {
-        return d1 == d2;
+      .def("__eq__", [](const Device& d, const nb::object& other) {
+        if (!nb::isinstance<Device>(other) &&
+            !nb::isinstance<Device::DeviceType>(other)) {
+          return false;
+        }
+        return d == nb::cast<Device>(other);
       });
 
-  py::implicitly_convertible<Device::DeviceType, Device>();
+  nb::implicitly_convertible<Device::DeviceType, Device>();
 
   m.def(
       "default_device",

python/src/fast.cpp

@@ -1,23 +1,87 @@
 // Copyright © 2023-2024 Apple Inc.
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/fast.h"
 #include "mlx/ops.h"
-#include "python/src/utils.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 
-void init_extensions(py::module_& parent_module) {
-  py::options options;
-  options.disable_function_signatures();
-
+void init_fast(nb::module_& parent_module) {
   auto m =
       parent_module.def_submodule("fast", "mlx.core.fast: fast operations");
 
+  m.def(
+      "rms_norm",
+      [](const array& x,
+         const array& weight,
+         float eps,
+         const StreamOrDevice& s /* = {} */) {
+        return fast::rms_norm(x, weight, eps, s);
+      },
+      "x"_a,
+      "weight"_a,
+      "eps"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def rms_norm(x: array, weight: array, eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Root Mean Square normalization (RMS norm).
+
+        The normalization is with respect to the last axis of the input ``x``.
+
+        Args:
+            x (array): Input array.
+            weight (array): A multiplicative weight to scale the result by.
+              The ``weight`` should be one-dimensional with the same size
+              as the last axis of ``x``.
+            eps (float): A small additive constant for numerical stability.
+
+        Returns:
+            array: The output array.
+      )pbdoc");
+
+  m.def(
+      "layer_norm",
+      [](const array& x,
+         const std::optional<array>& weight,
+         const std::optional<array>& bias,
+         float eps,
+         const StreamOrDevice& s /* = {} */) {
+        return fast::layer_norm(x, weight, bias, eps, s);
+      },
+      "x"_a,
+      "weight"_a.none(),
+      "bias"_a.none(),
+      "eps"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Layer normalization.
+
+        The normalization is with respect to the last axis of the input ``x``.
+
+        Args:
+            x (array): Input array.
+            weight (array, optional): A multiplicative weight to scale the result by.
+              The ``weight`` should be one-dimensional with the same size
+              as the last axis of ``x``. If set to ``None`` then no scaling happens.
+            bias (array, optional): An additive offset to be added to the result.
+              The ``bias`` should be one-dimensional with the same size
+              as the last axis of ``x``. If set to ``None`` then no translation happens.
+            eps (float): A small additive constant for numerical stability.
+
+        Returns:
+            array: The output array.
+      )pbdoc");
+
   m.def(
       "rope",
       [](const array& a,
@@ -31,15 +95,15 @@ void init_extensions(py::module_& parent_module) {
       },
       "a"_a,
       "dims"_a,
-      py::kw_only(),
+      nb::kw_only(),
       "traditional"_a,
       "base"_a,
       "scale"_a,
       "offset"_a,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def rope(a: array, dims: int, *, traditinoal: bool, base: float, scale: float, offset: int, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
-        rope(a: array, dims: int, *, traditinoal: bool, base: float, scale: float, offset: int, stream: Union[None, Stream, Device] = None) -> array
-
         Apply rotary positional encoding to the input.
 
         Args:
@@ -70,30 +134,34 @@ void init_extensions(py::module_& parent_module) {
       "q"_a,
       "k"_a,
       "v"_a,
-      py::kw_only(),
+      nb::kw_only(),
       "scale"_a,
-      "mask"_a = none,
-      "stream"_a = none,
+      "mask"_a = nb::none(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: Union[None, array] = None, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
-                  scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: Union[None, array] = None, stream: Union[None, Stream, Device] = None) -> array
+        A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``.
 
-            A fast implementation of multi-head attention: O = softmax(Q @ K.T, dim=-1) @ V.
-            Supports [Multi-Head Attention](https://arxiv.org/abs/1706.03762), [Grouped Query Attention](https://arxiv.org/abs/2305.13245), and [Multi-Query Attention](https://arxiv.org/abs/1911.02150).
+        Supports:
+        * [Multi-Head Attention](https://arxiv.org/abs/1706.03762)
+        * [Grouped Query Attention](https://arxiv.org/abs/2305.13245)
+        * [Multi-Query Attention](https://arxiv.org/abs/1911.02150).
 
-            This function will dispatch to an optimized Metal kernel when the query sequence length is 1. It handles other cases with regular MLX operations.
+        Note: The softmax operation is performed in ``float32`` regardless of
+        input precision.
 
-            Note: The softmax operation is performed in float32 precision regardless of input precision (float16 or float32).
-            Note: For Grouped Query Attention and Multi-Query Attention, the input arrays for `key` and `value` should not be pre-tiled to match the `query` array.
+        Note: For Grouped Query Attention and Multi-Query Attention, the ``k``
+        and ``v`` inputs should not be pre-tiled to match ``q``.
 
-            Args:
-                q (array): Input query array.
-                k (array): Input keys array.
-                v (array): Input values array.
-                scale (float): Scale for queries (typically ``1.0 / sqrt(q.shape(-1)``)
-                mask (array, optional): An additive mask to apply to the query-key scores.
+        Args:
+            q (array): Input query array.
+            k (array): Input keys array.
+            v (array): Input values array.
+            scale (float): Scale for queries (typically ``1.0 / sqrt(q.shape(-1)``)
+            mask (array, optional): An additive mask to apply to the query-key scores.
 
-            Returns:
-                array: The output array.
-
-          )pbdoc");
+        Returns:
+            array: The output array.
+      )pbdoc");
 }

python/src/fft.cpp

@@ -1,19 +1,20 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
-
-#include "python/src/utils.h"
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/variant.h>
+#include <nanobind/stl/vector.h>
+#include <numeric>
 
 #include "mlx/fft.h"
 #include "mlx/ops.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 
 using namespace mlx::core;
 
-void init_fft(py::module_& parent_module) {
+void init_fft(nb::module_& parent_module) {
   auto m = parent_module.def_submodule(
       "fft", "mlx.core.fft: Fast Fourier Transforms.");
   m.def(
@@ -29,9 +30,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "n"_a = none,
+      "n"_a = nb::none(),
       "axis"_a = -1,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
       R"pbdoc(
         One dimensional discrete Fourier Transform.
 
@@ -59,9 +60,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "n"_a = none,
+      "n"_a = nb::none(),
       "axis"_a = -1,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
       R"pbdoc(
         One dimensional inverse discrete Fourier Transform.
 
@@ -95,9 +96,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = std::vector<int>{-2, -1},
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a.none() = std::vector<int>{-2, -1},
+      "stream"_a = nb::none(),
       R"pbdoc(
         Two dimensional discrete Fourier Transform.
 
@@ -132,9 +133,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = std::vector<int>{-2, -1},
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a.none() = std::vector<int>{-2, -1},
+      "stream"_a = nb::none(),
       R"pbdoc(
         Two dimensional inverse discrete Fourier Transform.
 
@@ -169,9 +170,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = none,
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         n-dimensional discrete Fourier Transform.
 
@@ -207,9 +208,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = none,
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         n-dimensional inverse discrete Fourier Transform.
 
@@ -239,9 +240,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "n"_a = none,
+      "n"_a = nb::none(),
       "axis"_a = -1,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
       R"pbdoc(
         One dimensional discrete Fourier Transform on a real input.
 
@@ -274,9 +275,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "n"_a = none,
+      "n"_a = nb::none(),
       "axis"_a = -1,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
       R"pbdoc(
         The inverse of :func:`rfft`.
 
@@ -314,9 +315,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = std::vector<int>{-2, -1},
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a.none() = std::vector<int>{-2, -1},
+      "stream"_a = nb::none(),
       R"pbdoc(
         Two dimensional real discrete Fourier Transform.
 
@@ -357,9 +358,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = std::vector<int>{-2, -1},
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a.none() = std::vector<int>{-2, -1},
+      "stream"_a = nb::none(),
       R"pbdoc(
         The inverse of :func:`rfft2`.
 
@@ -400,9 +401,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = none,
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         n-dimensional real discrete Fourier Transform.
 
@@ -443,9 +444,9 @@ void init_fft(py::module_& parent_module) {
         }
       },
       "a"_a,
-      "s"_a = none,
-      "axes"_a = none,
-      "stream"_a = none,
+      "s"_a = nb::none(),
+      "axes"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         The inverse of :func:`rfftn`.
 

python/src/indexing.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <numeric>
 #include <sstream>
@@ -7,19 +7,19 @@
 
 #include "mlx/ops.h"
 
-bool is_none_slice(const py::slice& in_slice) {
+bool is_none_slice(const nb::slice& in_slice) {
   return (
-      py::getattr(in_slice, "start").is_none() &&
-      py::getattr(in_slice, "stop").is_none() &&
-      py::getattr(in_slice, "step").is_none());
+      nb::getattr(in_slice, "start").is_none() &&
+      nb::getattr(in_slice, "stop").is_none() &&
+      nb::getattr(in_slice, "step").is_none());
 }
 
-int get_slice_int(py::object obj, int default_val) {
+int get_slice_int(nb::object obj, int default_val) {
   if (!obj.is_none()) {
-    if (!py::isinstance<py::int_>(obj)) {
+    if (!nb::isinstance<nb::int_>(obj)) {
       throw std::invalid_argument("Slice indices must be integers or None.");
     }
-    return py::cast<int>(py::cast<py::int_>(obj));
+    return nb::cast<int>(nb::cast<nb::int_>(obj));
   }
   return default_val;
 }
@@ -28,7 +28,7 @@ void get_slice_params(
     int& starts,
     int& ends,
     int& strides,
-    const py::slice& in_slice,
+    const nb::slice& in_slice,
     int axis_size) {
   // Following numpy's convention
   //    Assume n is the number of elements in the dimension being sliced.
@@ -36,26 +36,26 @@ void get_slice_params(
   //    k < 0 . If j is not given it defaults to n for k > 0 and -n-1 for
   //    k < 0 . If k is not given it defaults to 1
 
-  strides = get_slice_int(py::getattr(in_slice, "step"), 1);
+  strides = get_slice_int(nb::getattr(in_slice, "step"), 1);
   starts = get_slice_int(
-      py::getattr(in_slice, "start"), strides < 0 ? axis_size - 1 : 0);
+      nb::getattr(in_slice, "start"), strides < 0 ? axis_size - 1 : 0);
   ends = get_slice_int(
-      py::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
+      nb::getattr(in_slice, "stop"), strides < 0 ? -axis_size - 1 : axis_size);
 }
 
-array get_int_index(py::object idx, int axis_size) {
-  int idx_ = py::cast<int>(idx);
+array get_int_index(nb::object idx, int axis_size) {
+  int idx_ = nb::cast<int>(idx);
   idx_ = (idx_ < 0) ? idx_ + axis_size : idx_;
 
   return array(idx_, uint32);
 }
 
-bool is_valid_index_type(const py::object& obj) {
-  return py::isinstance<py::slice>(obj) || py::isinstance<py::int_>(obj) ||
-      py::isinstance<array>(obj) || obj.is_none() || py::ellipsis().is(obj);
+bool is_valid_index_type(const nb::object& obj) {
+  return nb::isinstance<nb::slice>(obj) || nb::isinstance<nb::int_>(obj) ||
+      nb::isinstance<array>(obj) || obj.is_none() || nb::ellipsis().is(obj);
 }
 
-array mlx_get_item_slice(const array& src, const py::slice& in_slice) {
+array mlx_get_item_slice(const array& src, const nb::slice& in_slice) {
   // Check input and raise error if 0 dim for parity with np
   if (src.ndim() == 0) {
     throw std::invalid_argument(
@@ -92,7 +92,7 @@ array mlx_get_item_array(const array& src, const array& indices) {
   return take(src, indices, 0);
 }
 
-array mlx_get_item_int(const array& src, const py::int_& idx) {
+array mlx_get_item_int(const array& src, const nb::int_& idx) {
   // Check input and raise error if 0 dim for parity with np
   if (src.ndim() == 0) {
     throw std::invalid_argument(
@@ -106,7 +106,7 @@ array mlx_get_item_int(const array& src, const py::int_& idx) {
 
 array mlx_gather_nd(
     array src,
-    const std::vector<py::object>& indices,
+    const std::vector<nb::object>& indices,
     bool gather_first,
     int& max_dims) {
   max_dims = 0;
@@ -117,9 +117,10 @@ array mlx_gather_nd(
   for (int i = 0; i < indices.size(); i++) {
     auto& idx = indices[i];
 
-    if (py::isinstance<py::slice>(idx)) {
+    if (nb::isinstance<nb::slice>(idx)) {
       int start, end, stride;
-      get_slice_params(start, end, stride, idx, src.shape(i));
+      get_slice_params(
+          start, end, stride, nb::cast<nb::slice>(idx), src.shape(i));
 
       // Handle negative indices
       start = (start < 0) ? start + src.shape(i) : start;
@@ -128,10 +129,10 @@ array mlx_gather_nd(
       gather_indices.push_back(arange(start, end, stride, uint32));
       num_slices++;
       is_slice[i] = true;
-    } else if (py::isinstance<py::int_>(idx)) {
+    } else if (nb::isinstance<nb::int_>(idx)) {
       gather_indices.push_back(get_int_index(idx, src.shape(i)));
-    } else if (py::isinstance<array>(idx)) {
-      auto arr = py::cast<array>(idx);
+    } else if (nb::isinstance<array>(idx)) {
+      auto arr = nb::cast<array>(idx);
       max_dims = std::max(static_cast<int>(arr.ndim()), max_dims);
       gather_indices.push_back(arr);
     }
@@ -185,7 +186,65 @@ array mlx_gather_nd(
   return src;
 }
 
-array mlx_get_item_nd(array src, const py::tuple& entries) {
+auto mlx_expand_ellipsis(
+    const std::vector<int>& shape,
+    const nb::tuple& entries) {
+  std::vector<nb::object> indices;
+
+  // Go over all entries and note the position of ellipsis
+  int non_none_indices_before = 0;
+  int non_none_indices_after = 0;
+  std::vector<nb::object> r_indices;
+  int i = 0;
+
+  // Start from dimension 0 till we hit an ellipsis
+  for (; i < entries.size(); i++) {
+    auto idx = entries[i];
+    if (!is_valid_index_type(idx)) {
+      throw std::invalid_argument(
+          "Cannot index mlx array using the given type yet");
+    }
+    if (!nb::ellipsis().is(idx)) {
+      indices.push_back(idx);
+      non_none_indices_before += !idx.is_none();
+    } else {
+      break;
+    }
+  }
+
+  // If we do hit an ellipsis, collect indices from the back
+  for (int j = entries.size() - 1; j > i; j--) {
+    auto idx = entries[j];
+    if (!is_valid_index_type(idx)) {
+      throw std::invalid_argument(
+          "Cannot index mlx array using the given type yet");
+    }
+    if (nb::ellipsis().is(idx)) {
+      throw std::invalid_argument(
+          "An index can only have a single ellipsis (...)");
+    }
+    r_indices.push_back(idx);
+    non_none_indices_after += !idx.is_none();
+  }
+
+  // Count up the number of non none indices
+  int non_none_indices = non_none_indices_before + non_none_indices_after;
+
+  // Expand ellipsis
+  for (int axis = non_none_indices_before;
+       axis < shape.size() - non_none_indices_after;
+       axis++) {
+    indices.push_back(nb::slice(0, shape[axis], 1));
+    non_none_indices++;
+  }
+
+  // Insert indices collected after the ellipsis
+  indices.insert(indices.end(), r_indices.rbegin(), r_indices.rend());
+
+  return std::make_pair(non_none_indices, indices);
+}
+
+array mlx_get_item_nd(array src, const nb::tuple& entries) {
   // No indices make this a noop
   if (entries.size() == 0) {
     return src;
@@ -197,66 +256,20 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
   // 3. Calculate the remaining slices and reshapes
 
   // Ellipsis handling
-  std::vector<py::object> indices;
-  {
-    int non_none_indices_before = 0;
-    int non_none_indices_after = 0;
-    std::vector<py::object> r_indices;
-    int i = 0;
-    for (; i < entries.size(); i++) {
-      auto idx = entries[i];
-      if (!is_valid_index_type(idx)) {
-        throw std::invalid_argument(
-            "Cannot index mlx array using the given type yet");
-      }
-      if (!py::ellipsis().is(idx)) {
-        indices.push_back(idx);
-        non_none_indices_before += !idx.is_none();
-      } else {
-        break;
-      }
-    }
-    for (int j = entries.size() - 1; j > i; j--) {
-      auto idx = entries[j];
-      if (!is_valid_index_type(idx)) {
-        throw std::invalid_argument(
-            "Cannot index mlx array using the given type yet");
-      }
-      if (py::ellipsis().is(idx)) {
-        throw std::invalid_argument(
-            "An index can only have a single ellipsis (...)");
-      }
-      r_indices.push_back(idx);
-      non_none_indices_after += !idx.is_none();
-    }
-    for (int axis = non_none_indices_before;
-         axis < src.ndim() - non_none_indices_after;
-         axis++) {
-      indices.push_back(py::slice(0, src.shape(axis), 1));
-    }
-    indices.insert(indices.end(), r_indices.rbegin(), r_indices.rend());
-  }
+  auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
 
   // Check for the number of indices passed
-  {
-    int cnt = src.ndim();
-    for (auto& idx : indices) {
-      if (!idx.is_none()) {
-        cnt--;
-      }
-    }
-    if (cnt < 0) {
-      std::ostringstream msg;
-      msg << "Too many indices for array with " << src.ndim() << "dimensions.";
-      throw std::invalid_argument(msg.str());
-    }
+  if (non_none_indices > src.ndim()) {
+    std::ostringstream msg;
+    msg << "Too many indices for array with " << src.ndim() << "dimensions.";
+    throw std::invalid_argument(msg.str());
   }
 
   // Gather handling
   //
   // Check whether we have arrays or integer indices and delegate to gather_nd
   // after removing the slices at the end and all Nones.
-  std::vector<py::object> remaining_indices;
+  std::vector<nb::object> remaining_indices;
   bool have_array = false;
   {
     // First check whether the results of gather are going to be 1st or
@@ -264,7 +277,7 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
     bool have_non_array = false;
     bool gather_first = false;
     for (auto& idx : indices) {
-      if (py::isinstance<array>(idx) || py::isinstance<py::int_>(idx)) {
+      if (nb::isinstance<array>(idx) || (nb::isinstance<nb::int_>(idx))) {
         if (have_array && have_non_array) {
           gather_first = true;
           break;
@@ -275,17 +288,24 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
       }
     }
 
+    int n_arr = 0;
+    for (auto& idx : indices) {
+      n_arr += nb::isinstance<array>(idx);
+    }
+
+    have_array &= n_arr > 0;
+
     if (have_array) {
       int last_array;
       // Then find the last array
       for (last_array = indices.size() - 1; last_array >= 0; last_array--) {
         auto& idx = indices[last_array];
-        if (py::isinstance<array>(idx) || py::isinstance<py::int_>(idx)) {
+        if (nb::isinstance<array>(idx) || nb::isinstance<nb::int_>(idx)) {
           break;
         }
       }
 
-      std::vector<py::object> gather_indices;
+      std::vector<nb::object> gather_indices;
       for (int i = 0; i <= last_array; i++) {
         auto& idx = indices[i];
         if (!idx.is_none()) {
@@ -299,15 +319,15 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
       if (gather_first) {
         for (int i = 0; i < max_dims; i++) {
           remaining_indices.push_back(
-              py::slice(py::none(), py::none(), py::none()));
+              nb::slice(nb::none(), nb::none(), nb::none()));
         }
         for (int i = 0; i < last_array; i++) {
           auto& idx = indices[i];
           if (idx.is_none()) {
             remaining_indices.push_back(indices[i]);
-          } else if (py::isinstance<py::slice>(idx)) {
+          } else if (nb::isinstance<nb::slice>(idx)) {
             remaining_indices.push_back(
-                py::slice(py::none(), py::none(), py::none()));
+                nb::slice(nb::none(), nb::none(), nb::none()));
           }
         }
         for (int i = last_array + 1; i < indices.size(); i++) {
@@ -316,18 +336,18 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
       } else {
         for (int i = 0; i < indices.size(); i++) {
           auto& idx = indices[i];
-          if (py::isinstance<array>(idx) || py::isinstance<py::int_>(idx)) {
+          if (nb::isinstance<array>(idx) || nb::isinstance<nb::int_>(idx)) {
             break;
           } else if (idx.is_none()) {
             remaining_indices.push_back(idx);
           } else {
             remaining_indices.push_back(
-                py::slice(py::none(), py::none(), py::none()));
+                nb::slice(nb::none(), nb::none(), nb::none()));
           }
         }
         for (int i = 0; i < max_dims; i++) {
           remaining_indices.push_back(
-              py::slice(py::none(), py::none(), py::none()));
+              nb::slice(nb::none(), nb::none(), nb::none()));
         }
         for (int i = last_array + 1; i < indices.size(); i++) {
           remaining_indices.push_back(indices[i]);
@@ -342,6 +362,8 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
     remaining_indices = indices;
   }
 
+  bool squeeze_needed = false;
+
   // Slice handling
   {
     std::vector<int> starts(src.ndim(), 0);
@@ -350,8 +372,24 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
     int axis = 0;
     for (auto& idx : remaining_indices) {
       if (!idx.is_none()) {
-        get_slice_params(
-            starts[axis], ends[axis], strides[axis], idx, ends[axis]);
+        if (!have_array && nb::isinstance<nb::int_>(idx)) {
+          int st = nb::cast<int>(idx);
+          st = (st < 0) ? st + src.shape(axis) : st;
+
+          starts[axis] = st;
+          ends[axis] = st + 1;
+
+          squeeze_needed = true;
+
+        } else {
+          get_slice_params(
+              starts[axis],
+              ends[axis],
+              strides[axis],
+              nb::cast<nb::slice>(idx),
+              ends[axis]);
+        }
+
         axis++;
       }
     }
@@ -359,12 +397,14 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
   }
 
   // Unsqueeze handling
-  if (remaining_indices.size() > src.ndim()) {
+  if (remaining_indices.size() > src.ndim() || squeeze_needed) {
     std::vector<int> out_shape;
     int axis = 0;
     for (auto& idx : remaining_indices) {
       if (idx.is_none()) {
         out_shape.push_back(1);
+      } else if (squeeze_needed && nb::isinstance<nb::int_>(idx)) {
+        axis++;
       } else {
         out_shape.push_back(src.shape(axis++));
       }
@@ -375,15 +415,17 @@ array mlx_get_item_nd(array src, const py::tuple& entries) {
   return src;
 }
 
-array mlx_get_item(const array& src, const py::object& obj) {
-  if (py::isinstance<py::slice>(obj)) {
-    return mlx_get_item_slice(src, obj);
-  } else if (py::isinstance<array>(obj)) {
-    return mlx_get_item_array(src, py::cast<array>(obj));
-  } else if (py::isinstance<py::int_>(obj)) {
-    return mlx_get_item_int(src, obj);
-  } else if (py::isinstance<py::tuple>(obj)) {
-    return mlx_get_item_nd(src, obj);
+array mlx_get_item(const array& src, const nb::object& obj) {
+  if (nb::isinstance<nb::slice>(obj)) {
+    return mlx_get_item_slice(src, nb::cast<nb::slice>(obj));
+  } else if (nb::isinstance<array>(obj)) {
+    return mlx_get_item_array(src, nb::cast<array>(obj));
+  } else if (nb::isinstance<nb::int_>(obj)) {
+    return mlx_get_item_int(src, nb::cast<nb::int_>(obj));
+  } else if (nb::isinstance<nb::tuple>(obj)) {
+    return mlx_get_item_nd(src, nb::cast<nb::tuple>(obj));
+  } else if (nb::isinstance<nb::ellipsis>(obj)) {
+    return src;
   } else if (obj.is_none()) {
     std::vector<int> s(1, 1);
     s.insert(s.end(), src.shape().begin(), src.shape().end());
@@ -394,7 +436,7 @@ array mlx_get_item(const array& src, const py::object& obj) {
 
 std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_int(
     const array& src,
-    const py::int_& idx,
+    const nb::int_& idx,
     const array& update) {
   if (src.ndim() == 0) {
     throw std::invalid_argument(
@@ -446,7 +488,7 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_array(
 
 std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_slice(
     const array& src,
-    const py::slice& in_slice,
+    const nb::slice& in_slice,
     const array& update) {
   // Check input and raise error if 0 dim for parity with np
   if (src.ndim() == 0) {
@@ -472,55 +514,45 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_slice(
   // Check and update slice params
   get_slice_params(start, end, stride, in_slice, end);
 
+  // If simple stride
+  if (stride == 1) {
+    // Squeeze out singleton dims from the start of update
+    int s = 0;
+    for (; s < update.ndim() && update.shape(s) == 1; s++)
+      ;
+    auto up_shape =
+        std::vector<int>(update.shape().begin() + s, update.shape().end());
+    auto up = reshape(update, up_shape);
+
+    // Build array to mark start of slice
+    auto idx = array({start}, {1}, uint32);
+
+    // Get slice size
+    int slice_size = (end - start);
+
+    // Broadcast update to slide size
+    std::vector<int> up_shape_broadcast = {1, slice_size};
+    up_shape_broadcast.insert(
+        up_shape_broadcast.end(), src.shape().begin() + 1, src.shape().end());
+
+    up = broadcast_to(update, up_shape_broadcast);
+
+    auto indices = std::vector<array>{idx};
+    auto axes = std::vector<int>{0};
+
+    return {indices, up, axes};
+  }
+
   return mlx_scatter_args_array(
       src, arange(start, end, stride, uint32), update);
 }
 
 std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
     const array& src,
-    const py::tuple& entries,
+    const nb::tuple& entries,
     const array& update) {
-  std::vector<py::object> indices;
-  int non_none_indices = 0;
-
   // Expand ellipses into a series of ':' slices
-  {
-    int non_none_indices_before = 0;
-    int non_none_indices_after = 0;
-    bool has_ellipsis = false;
-    int indices_before = 0;
-    for (int i = 0; i < entries.size(); ++i) {
-      auto idx = entries[i];
-      if (!is_valid_index_type(idx)) {
-        throw std::invalid_argument(
-            "Cannot index mlx array using the given type yet");
-      } else if (!py::ellipsis().is(idx)) {
-        if (!has_ellipsis) {
-          indices_before++;
-          non_none_indices_before += !idx.is_none();
-        } else {
-          non_none_indices_after += !idx.is_none();
-        }
-        indices.push_back(idx);
-      } else if (has_ellipsis) {
-        throw std::invalid_argument(
-            "An index can only have a single ellipsis (...)");
-      } else {
-        has_ellipsis = true;
-      }
-    }
-    if (has_ellipsis) {
-      for (int axis = non_none_indices_before;
-           axis < src.ndim() - non_none_indices_after;
-           axis++) {
-        indices.insert(
-            indices.begin() + indices_before, py::slice(0, src.shape(axis), 1));
-      }
-      non_none_indices = src.ndim();
-    } else {
-      non_none_indices = non_none_indices_before + non_none_indices_after;
-    }
-  }
+  auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
 
   if (non_none_indices > src.ndim()) {
     std::ostringstream msg;
@@ -541,85 +573,142 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
     return {{}, broadcast_to(up, src.shape()), {}};
   }
 
+  // Analyse the types of the indices
   unsigned long max_dim = 0;
   bool arrays_first = false;
+  int num_none = 0;
   int num_slices = 0;
   int num_arrays = 0;
+  int num_strided_slices = 0;
   {
     bool have_array = false;
     bool have_non_array = false;
     for (auto& idx : indices) {
-      if (py::isinstance<py::slice>(idx) || idx.is_none()) {
+      if (idx.is_none()) {
+        have_non_array = have_array;
+        num_none++;
+
+      } else if (nb::isinstance<nb::slice>(idx)) {
         have_non_array = have_array;
         num_slices++;
-      } else if (py::isinstance<array>(idx)) {
+
+        auto slice = nb::cast<nb::slice>(idx);
+        int stride = get_slice_int(nb::getattr(slice, "step"), 1);
+        num_strided_slices += (stride != 1);
+
+      } else if (nb::isinstance<array>(idx)) {
         have_array = true;
         if (have_array && have_non_array) {
           arrays_first = true;
         }
-        max_dim = std::max(py::cast<array>(idx).ndim(), max_dim);
+        max_dim = std::max(nb::cast<array>(idx).ndim(), max_dim);
         num_arrays++;
       }
     }
   }
 
+  // We have index dims for the arrays, strided slices (implemented as arrays),
+  // none
+  int idx_ndim = max_dim + num_strided_slices + num_none;
+
+  // If we have simple non-strided slices, we also attach an index for that
+  idx_ndim += (num_slices < num_strided_slices);
+
+  // Go over each index type and translate to the needed scatter args
   std::vector<array> arr_indices;
   int slice_num = 0;
   int array_num = 0;
   int ax = 0;
+
+  // We collect the shapes of the slices and updates during this process
+  std::vector<int> update_shape(non_none_indices, 1);
+  std::vector<int> slice_shapes;
+
   for (int i = 0; i < indices.size(); ++i) {
     auto& pyidx = indices[i];
-    if (py::isinstance<py::slice>(pyidx)) {
+    if (nb::isinstance<nb::slice>(pyidx)) {
       int start, end, stride;
       auto axis_size = src.shape(ax++);
-      get_slice_params(start, end, stride, pyidx, axis_size);
+      get_slice_params(
+          start, end, stride, nb::cast<nb::slice>(pyidx), axis_size);
 
       // Handle negative indices
       start = (start < 0) ? start + axis_size : start;
       end = (end < 0) ? end + axis_size : end;
 
-      auto idx = arange(start, end, stride, uint32);
-      std::vector<int> idx_shape(max_dim + num_slices, 1);
-      auto loc = slice_num + (arrays_first ? max_dim : 0);
-      slice_num++;
-      idx_shape[loc] = idx.size();
-      arr_indices.push_back(reshape(idx, idx_shape));
-    } else if (py::isinstance<py::int_>(pyidx)) {
-      arr_indices.push_back(get_int_index(pyidx, src.shape(ax++)));
-    } else if (pyidx.is_none()) {
-      slice_num++;
-    } else if (py::isinstance<array>(pyidx)) {
-      ax++;
-      auto idx = py::cast<array>(pyidx);
-      std::vector<int> idx_shape;
-      if (!arrays_first) {
-        idx_shape.insert(idx_shape.end(), slice_num, 1);
+      std::vector<int> idx_shape(idx_ndim, 1);
+
+      // If it's a simple slice, we only need to add the start index
+      if (stride == 1) {
+        auto idx = array({start}, idx_shape, uint32);
+        slice_shapes.push_back(end - start);
+        arr_indices.push_back(idx);
+      }
+      // Otherwise we expand the slice into indices using arange
+      else {
+        auto idx = arange(start, end, stride, uint32);
+        auto loc = slice_num + (arrays_first ? max_dim : 0);
+        slice_num++;
+        idx_shape[loc] = idx.size();
+        slice_shapes.push_back(idx.size());
+        arr_indices.push_back(reshape(idx, idx_shape));
+      }
+      // Add the shape to the update
+      update_shape[ax - 1] = slice_shapes.back();
+    } else if (nb::isinstance<nb::int_>(pyidx)) {
+      // Add index to arrays
+      arr_indices.push_back(get_int_index(pyidx, src.shape(ax++)));
+      // Add the shape to the update
+      update_shape[ax - 1] = 1;
+    } else if (pyidx.is_none()) {
+      // We only use the None's for bookeeping dimensions
+      slice_num++;
+    } else if (nb::isinstance<array>(pyidx)) {
+      ax++;
+      auto idx = nb::cast<array>(pyidx);
+      std::vector<int> idx_shape(idx_ndim, 1);
+
+      // Place the arrays in the correct dimension
+      int st = (!arrays_first) * slice_num + max_dim - idx.ndim();
+      for (int j = 0; j < idx.ndim(); j++) {
+        idx_shape[st + j] = idx.shape()[j];
       }
-      idx_shape.insert(idx_shape.end(), max_dim - idx.ndim(), 1);
-      idx_shape.insert(idx_shape.end(), idx.shape().begin(), idx.shape().end());
-      idx_shape.insert(
-          idx_shape.end(), num_slices - (arrays_first ? 0 : slice_num), 1);
       arr_indices.push_back(reshape(idx, idx_shape));
       if (!arrays_first && ++array_num == num_arrays) {
         slice_num += max_dim;
       }
+
+      // Add the shape to the update
+      update_shape[ax - 1] = 1;
     } else {
       throw std::invalid_argument(
           "Cannot index mlx array using the given type yet");
     }
   }
 
+  // Broadcast the update to the indices and slices
   arr_indices = broadcast_arrays(arr_indices);
-  up_shape = arr_indices[0].shape();
-  up_shape.insert(
-      up_shape.end(),
+  auto up_shape_broadcast = arr_indices[0].shape();
+
+  up_shape_broadcast.insert(
+      up_shape_broadcast.end(), slice_shapes.begin(), slice_shapes.end());
+  up_shape_broadcast.insert(
+      up_shape_broadcast.end(),
       src.shape().begin() + non_none_indices,
       src.shape().end());
-  up = broadcast_to(up, up_shape);
-  up_shape.insert(
-      up_shape.begin() + arr_indices[0].ndim(), non_none_indices, 1);
-  up = reshape(up, up_shape);
+  up = broadcast_to(up, up_shape_broadcast);
 
+  // Reshape the update with the size-1 dims for the int and array indices
+  auto up_reshape = arr_indices[0].shape();
+  up_reshape.insert(up_reshape.end(), update_shape.begin(), update_shape.end());
+  up_reshape.insert(
+      up_reshape.end(),
+      src.shape().begin() + non_none_indices,
+      src.shape().end());
+
+  up = reshape(up, up_reshape);
+
+  // Collect axes
   std::vector<int> axes(arr_indices.size(), 0);
   std::iota(axes.begin(), axes.end(), 0);
 
@@ -629,24 +718,129 @@ std::tuple<std::vector<array>, array, std::vector<int>> mlx_scatter_args_nd(
 std::tuple<std::vector<array>, array, std::vector<int>>
 mlx_compute_scatter_args(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto vals = to_array(v, src.dtype());
-  if (py::isinstance<py::slice>(obj)) {
-    return mlx_scatter_args_slice(src, obj, vals);
-  } else if (py::isinstance<array>(obj)) {
-    return mlx_scatter_args_array(src, py::cast<array>(obj), vals);
-  } else if (py::isinstance<py::int_>(obj)) {
-    return mlx_scatter_args_int(src, obj, vals);
-  } else if (py::isinstance<py::tuple>(obj)) {
-    return mlx_scatter_args_nd(src, obj, vals);
+  if (nb::isinstance<nb::slice>(obj)) {
+    return mlx_scatter_args_slice(src, nb::cast<nb::slice>(obj), vals);
+  } else if (nb::isinstance<array>(obj)) {
+    return mlx_scatter_args_array(src, nb::cast<array>(obj), vals);
+  } else if (nb::isinstance<nb::int_>(obj)) {
+    return mlx_scatter_args_int(src, nb::cast<nb::int_>(obj), vals);
+  } else if (nb::isinstance<nb::tuple>(obj)) {
+    return mlx_scatter_args_nd(src, nb::cast<nb::tuple>(obj), vals);
   } else if (obj.is_none()) {
     return {{}, broadcast_to(vals, src.shape()), {}};
   }
   throw std::invalid_argument("Cannot index mlx array using the given type.");
 }
 
-void mlx_set_item(array& src, const py::object& obj, const ScalarOrArray& v) {
+auto mlx_slice_update(
+    const array& src,
+    const nb::object& obj,
+    const ScalarOrArray& v) {
+  // Can't route to slice update if not slice or tuple
+  if (src.ndim() == 0 ||
+      (!nb::isinstance<nb::slice>(obj) && !nb::isinstance<nb::tuple>(obj))) {
+    return std::make_pair(false, src);
+  }
+
+  // Should be able to route to slice update
+
+  // Pre process tuple
+  auto upd = to_array(v, src.dtype());
+
+  // Remove leading singletons dimensions from the update
+  int s = 0;
+  for (; s < upd.ndim() && upd.shape(s) == 1; s++) {
+  };
+  auto up_shape = std::vector<int>(upd.shape().begin() + s, upd.shape().end());
+  up_shape = up_shape.empty() ? std::vector{1} : up_shape;
+  auto up = reshape(upd, up_shape);
+
+  // Build slice update params
+  std::vector<int> starts(src.ndim(), 0);
+  std::vector<int> stops = src.shape();
+  std::vector<int> strides(src.ndim(), 1);
+
+  // If it's just a simple slice, just do a slice update and return
+  if (nb::isinstance<nb::slice>(obj)) {
+    // Read slice arguments
+    get_slice_params(
+        starts[0],
+        stops[0],
+        strides[0],
+        nb::cast<nb::slice>(obj),
+        src.shape(0));
+
+    // Do slice update
+    auto out = slice_update(src, up, starts, stops, strides);
+    return std::make_pair(true, out);
+  }
+
+  // It must be a tuple
+  auto entries = nb::cast<nb::tuple>(obj);
+
+  // Can't route to slice update if any arrays are present
+  for (int i = 0; i < entries.size(); i++) {
+    auto idx = entries[i];
+    if (nb::isinstance<array>(idx)) {
+      return std::make_pair(false, src);
+    }
+  }
+
+  // Expand ellipses into a series of ':' slices
+  auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries);
+
+  // Dimension check
+  if (non_none_indices > src.ndim()) {
+    std::ostringstream msg;
+    msg << "Too many indices for array with " << src.ndim() << "dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // If no non-None indices return the broadcasted update
+  if (non_none_indices == 0) {
+    return std::make_pair(true, broadcast_to(up, src.shape()));
+  }
+
+  // Process entries
+  std::vector<int> upd_expand_dims;
+  int ax = 0;
+  for (int i = 0; i < indices.size(); ++i) {
+    auto& pyidx = indices[i];
+    if (nb::isinstance<nb::slice>(pyidx)) {
+      get_slice_params(
+          starts[ax],
+          stops[ax],
+          strides[ax],
+          nb::cast<nb::slice>(pyidx),
+          src.shape(ax));
+      ax++;
+    } else if (nb::isinstance<nb::int_>(pyidx)) {
+      int st = nb::cast<int>(pyidx);
+      st = (st < 0) ? st + src.shape(ax) : st;
+      starts[ax] = st;
+      stops[ax] = st + 1;
+      if (src.ndim() - ax < up.ndim()) {
+        upd_expand_dims.push_back(ax - src.ndim());
+      }
+      ax++;
+    }
+  }
+
+  up = expand_dims(up, upd_expand_dims);
+  auto out = slice_update(src, up, starts, stops, strides);
+  return std::make_pair(true, out);
+}
+
+void mlx_set_item(array& src, const nb::object& obj, const ScalarOrArray& v) {
+  auto [success, out] = mlx_slice_update(src, obj, v);
+  if (success) {
+    src.overwrite_descriptor(out);
+    return;
+  }
+
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
     auto out = scatter(src, indices, updates, axes);
@@ -658,7 +852,7 @@ void mlx_set_item(array& src, const py::object& obj, const ScalarOrArray& v) {
 
 array mlx_add_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
@@ -670,7 +864,7 @@ array mlx_add_item(
 
 array mlx_subtract_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
@@ -682,7 +876,7 @@ array mlx_subtract_item(
 
 array mlx_multiply_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
@@ -694,7 +888,7 @@ array mlx_multiply_item(
 
 array mlx_divide_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
@@ -706,7 +900,7 @@ array mlx_divide_item(
 
 array mlx_maximum_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {
@@ -718,7 +912,7 @@ array mlx_maximum_item(
 
 array mlx_minimum_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v) {
   auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v);
   if (indices.size() > 0) {

python/src/indexing.h

@@ -1,38 +1,38 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
 
 #include "mlx/array.h"
 #include "python/src/utils.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 using namespace mlx::core;
 
-array mlx_get_item(const array& src, const py::object& obj);
-void mlx_set_item(array& src, const py::object& obj, const ScalarOrArray& v);
+array mlx_get_item(const array& src, const nb::object& obj);
+void mlx_set_item(array& src, const nb::object& obj, const ScalarOrArray& v);
 array mlx_add_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);
 array mlx_subtract_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);
 array mlx_multiply_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);
 array mlx_divide_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);
 array mlx_maximum_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);
 array mlx_minimum_item(
     const array& src,
-    const py::object& obj,
+    const nb::object& obj,
     const ScalarOrArray& v);

python/src/linalg.cpp

@@ -1,32 +1,29 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <variant>
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/pair.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/variant.h>
+#include <nanobind/stl/vector.h>
 
 #include "mlx/linalg.h"
 
-#include "python/src/load.h"
-#include "python/src/utils.h"
-
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 
 using namespace mlx::core;
 using namespace mlx::core::linalg;
 
 namespace {
-py::tuple svd_helper(const array& a, StreamOrDevice s /* = {} */) {
+nb::tuple svd_helper(const array& a, StreamOrDevice s /* = {} */) {
   const auto result = svd(a, s);
-  return py::make_tuple(result.at(0), result.at(1), result.at(2));
+  return nb::make_tuple(result.at(0), result.at(1), result.at(2));
 }
 } // namespace
 
-void init_linalg(py::module_& parent_module) {
-  py::options options;
-  options.disable_function_signatures();
-
+void init_linalg(nb::module_& parent_module) {
   auto m = parent_module.def_submodule(
       "linalg", "mlx.core.linalg: linear algebra routines.");
 
@@ -59,16 +56,15 @@ void init_linalg(py::module_& parent_module) {
           return norm(a, ord, axis, keepdims, stream);
         }
       },
-      "a"_a,
-      py::pos_only(),
-      "ord"_a = none,
-      "axis"_a = none,
+      nb::arg(),
+      "ord"_a = nb::none(),
+      "axis"_a = nb::none(),
       "keepdims"_a = false,
-      py::kw_only(),
-      "stream"_a = none,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def norm(a: array, /, ord: Union[None, scalar, str] = None, axis: Union[None, int, List[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
-        norm(a: array, /, ord: Union[None, scalar, str] = None, axis: Union[None, int, List[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array
-
         Matrix or vector norm.
 
         This function computes vector or  matrix norms depending on the value of
@@ -188,11 +184,11 @@ void init_linalg(py::module_& parent_module) {
       "qr",
       &qr,
       "a"_a,
-      py::kw_only(),
-      "stream"_a = none,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def qr(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array)"),
       R"pbdoc(
-        qr(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array)
-
         The QR factorization of the input matrix.
 
         This function supports arrays with at least 2 dimensions. The matrices
@@ -221,11 +217,11 @@ void init_linalg(py::module_& parent_module) {
       "svd",
       &svd_helper,
       "a"_a,
-      py::kw_only(),
-      "stream"_a = none,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def svd(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array, array)"),
       R"pbdoc(
-        svd(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array, array)
-
         The Singular Value Decomposition (SVD) of the input matrix.
 
         This function supports arrays with at least 2 dimensions. When the input
@@ -241,4 +237,27 @@ void init_linalg(py::module_& parent_module) {
             tuple(array, array, array): The ``U``, ``S``, and ``Vt`` matrices, such that
               ``A = U @ diag(S) @ Vt``
       )pbdoc");
+  m.def(
+      "inv",
+      &inv,
+      "a"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def inv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the inverse of a square matrix.
+
+        This function supports arrays with at least 2 dimensions. When the input
+        has more than two dimensions, the inverse is computed for each matrix
+        in the last two dimensions of ``a``.
+
+        Args:
+            a (array): Input array.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+            array: ``ainv`` such that ``dot(a, ainv) = dot(ainv, a) = eye(a.shape[0])``
+      )pbdoc");
 }

python/src/load.cpp

@@ -1,8 +1,6 @@
-// Copyright © 2023 Apple Inc.
-
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+// Copyright © 2023-2024 Apple Inc.
 
+#include <nanobind/stl/vector.h>
 #include <cstring>
 #include <fstream>
 #include <stdexcept>
@@ -16,39 +14,39 @@
 #include "python/src/load.h"
 #include "python/src/utils.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 
 ///////////////////////////////////////////////////////////////////////////////
 // Helpers
 ///////////////////////////////////////////////////////////////////////////////
 
-bool is_istream_object(const py::object& file) {
-  return py::hasattr(file, "readinto") && py::hasattr(file, "seek") &&
-      py::hasattr(file, "tell") && py::hasattr(file, "closed");
+bool is_istream_object(const nb::object& file) {
+  return nb::hasattr(file, "readinto") && nb::hasattr(file, "seek") &&
+      nb::hasattr(file, "tell") && nb::hasattr(file, "closed");
 }
 
-bool is_ostream_object(const py::object& file) {
-  return py::hasattr(file, "write") && py::hasattr(file, "seek") &&
-      py::hasattr(file, "tell") && py::hasattr(file, "closed");
+bool is_ostream_object(const nb::object& file) {
+  return nb::hasattr(file, "write") && nb::hasattr(file, "seek") &&
+      nb::hasattr(file, "tell") && nb::hasattr(file, "closed");
 }
 
-bool is_zip_file(const py::module_& zipfile, const py::object& file) {
+bool is_zip_file(const nb::module_& zipfile, const nb::object& file) {
   if (is_istream_object(file)) {
     auto st_pos = file.attr("tell")();
-    bool r = (zipfile.attr("is_zipfile")(file)).cast<bool>();
+    bool r = nb::cast<bool>(zipfile.attr("is_zipfile")(file));
     file.attr("seek")(st_pos, 0);
     return r;
   }
-  return zipfile.attr("is_zipfile")(file).cast<bool>();
+  return nb::cast<bool>(zipfile.attr("is_zipfile")(file));
 }
 
 class ZipFileWrapper {
  public:
   ZipFileWrapper(
-      const py::module_& zipfile,
-      const py::object& file,
+      const nb::module_& zipfile,
+      const nb::object& file,
       char mode = 'r',
       int compression = 0)
       : zipfile_module_(zipfile),
@@ -63,10 +61,10 @@ class ZipFileWrapper {
         close_func_(zipfile_object_.attr("close")) {}
 
   std::vector<std::string> namelist() const {
-    return files_list_.cast<std::vector<std::string>>();
+    return nb::cast<std::vector<std::string>>(files_list_);
   }
 
-  py::object open(const std::string& key, char mode = 'r') {
+  nb::object open(const std::string& key, char mode = 'r') {
     // Following numpy :
     // https://github.com/numpy/numpy/blob/db4f43983cb938f12c311e1f5b7165e270c393b4/numpy/lib/npyio.py#L742C36-L742C47
     if (mode == 'w') {
@@ -76,12 +74,12 @@ class ZipFileWrapper {
   }
 
  private:
-  py::module_ zipfile_module_;
-  py::object zipfile_object_;
-  py::list files_list_;
-  py::object open_func_;
-  py::object read_func_;
-  py::object close_func_;
+  nb::module_ zipfile_module_;
+  nb::object zipfile_object_;
+  nb::list files_list_;
+  nb::object open_func_;
+  nb::object read_func_;
+  nb::object close_func_;
 };
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -90,14 +88,14 @@ class ZipFileWrapper {
 
 class PyFileReader : public io::Reader {
  public:
-  PyFileReader(py::object file)
+  PyFileReader(nb::object file)
       : pyistream_(file),
         readinto_func_(file.attr("readinto")),
         seek_func_(file.attr("seek")),
         tell_func_(file.attr("tell")) {}
 
   ~PyFileReader() {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
 
     pyistream_.release().dec_ref();
     readinto_func_.release().dec_ref();
@@ -108,8 +106,8 @@ class PyFileReader : public io::Reader {
   bool is_open() const override {
     bool out;
     {
-      py::gil_scoped_acquire gil;
-      out = !pyistream_.attr("closed").cast<bool>();
+      nb::gil_scoped_acquire gil;
+      out = !nb::cast<bool>(pyistream_.attr("closed"));
     }
     return out;
   }
@@ -117,7 +115,7 @@ class PyFileReader : public io::Reader {
   bool good() const override {
     bool out;
     {
-      py::gil_scoped_acquire gil;
+      nb::gil_scoped_acquire gil;
       out = !pyistream_.is_none();
     }
     return out;
@@ -126,25 +124,24 @@ class PyFileReader : public io::Reader {
   size_t tell() const override {
     size_t out;
     {
-      py::gil_scoped_acquire gil;
-      out = tell_func_().cast<size_t>();
+      nb::gil_scoped_acquire gil;
+      out = nb::cast<size_t>(tell_func_());
     }
     return out;
   }
 
   void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
       override {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
     seek_func_(off, (int)way);
   }
 
   void read(char* data, size_t n) override {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
+    auto memview = PyMemoryView_FromMemory(data, n, PyBUF_WRITE);
+    nb::object bytes_read = readinto_func_(nb::handle(memview));
 
-    py::object bytes_read =
-        readinto_func_(py::memoryview::from_buffer(data, {n}, {sizeof(char)}));
-
-    if (bytes_read.is_none() || py::cast<size_t>(bytes_read) < n) {
+    if (bytes_read.is_none() || nb::cast<size_t>(bytes_read) < n) {
       throw std::runtime_error("[load] Failed to read from python stream");
     }
   }
@@ -154,23 +151,23 @@ class PyFileReader : public io::Reader {
   }
 
  private:
-  py::object pyistream_;
-  py::object readinto_func_;
-  py::object seek_func_;
-  py::object tell_func_;
+  nb::object pyistream_;
+  nb::object readinto_func_;
+  nb::object seek_func_;
+  nb::object tell_func_;
 };
 
 std::pair<
     std::unordered_map<std::string, array>,
     std::unordered_map<std::string, std::string>>
-mlx_load_safetensor_helper(py::object file, StreamOrDevice s) {
-  if (py::isinstance<py::str>(file)) { // Assume .safetensors file path string
-    return load_safetensors(py::cast<std::string>(file), s);
+mlx_load_safetensor_helper(nb::object file, StreamOrDevice s) {
+  if (nb::isinstance<nb::str>(file)) { // Assume .safetensors file path string
+    return load_safetensors(nb::cast<std::string>(file), s);
   } else if (is_istream_object(file)) {
     // If we don't own the stream and it was passed to us, eval immediately
     auto res = load_safetensors(std::make_shared<PyFileReader>(file), s);
     {
-      py::gil_scoped_release gil;
+      nb::gil_scoped_release gil;
       for (auto& [key, arr] : std::get<0>(res)) {
         arr.eval();
       }
@@ -182,20 +179,20 @@ mlx_load_safetensor_helper(py::object file, StreamOrDevice s) {
       "[load_safetensors] Input must be a file-like object, or string");
 }
 
-GGUFLoad mlx_load_gguf_helper(py::object file, StreamOrDevice s) {
-  if (py::isinstance<py::str>(file)) { // Assume .gguf file path string
-    return load_gguf(py::cast<std::string>(file), s);
+GGUFLoad mlx_load_gguf_helper(nb::object file, StreamOrDevice s) {
+  if (nb::isinstance<nb::str>(file)) { // Assume .gguf file path string
+    return load_gguf(nb::cast<std::string>(file), s);
   }
 
   throw std::invalid_argument("[load_gguf] Input must be a string");
 }
 
 std::unordered_map<std::string, array> mlx_load_npz_helper(
-    py::object file,
+    nb::object file,
     StreamOrDevice s) {
-  bool own_file = py::isinstance<py::str>(file);
+  bool own_file = nb::isinstance<nb::str>(file);
 
-  py::module_ zipfile = py::module_::import("zipfile");
+  nb::module_ zipfile = nb::module_::import_("zipfile");
   if (!is_zip_file(zipfile, file)) {
     throw std::invalid_argument(
         "[load_npz] Input must be a zip file or a file-like object that can be "
@@ -208,9 +205,9 @@ std::unordered_map<std::string, array> mlx_load_npz_helper(
   ZipFileWrapper zipfile_object(zipfile, file);
   for (const std::string& st : zipfile_object.namelist()) {
     // Open zip file as a python file stream
-    py::object sub_file = zipfile_object.open(st);
+    nb::object sub_file = zipfile_object.open(st);
 
-    // Create array from python fille stream
+    // Create array from python file stream
     auto arr = load(std::make_shared<PyFileReader>(sub_file), s);
 
     // Remove .npy from file if it is there
@@ -224,7 +221,7 @@ std::unordered_map<std::string, array> mlx_load_npz_helper(
 
   // If we don't own the stream and it was passed to us, eval immediately
   if (!own_file) {
-    py::gil_scoped_release gil;
+    nb::gil_scoped_release gil;
     for (auto& [key, arr] : array_dict) {
       arr.eval();
     }
@@ -233,14 +230,14 @@ std::unordered_map<std::string, array> mlx_load_npz_helper(
   return array_dict;
 }
 
-array mlx_load_npy_helper(py::object file, StreamOrDevice s) {
-  if (py::isinstance<py::str>(file)) { // Assume .npy file path string
-    return load(py::cast<std::string>(file), s);
+array mlx_load_npy_helper(nb::object file, StreamOrDevice s) {
+  if (nb::isinstance<nb::str>(file)) { // Assume .npy file path string
+    return load(nb::cast<std::string>(file), s);
   } else if (is_istream_object(file)) {
     // If we don't own the stream and it was passed to us, eval immediately
     auto arr = load(std::make_shared<PyFileReader>(file), s);
     {
-      py::gil_scoped_release gil;
+      nb::gil_scoped_release gil;
       arr.eval();
     }
     return arr;
@@ -250,16 +247,16 @@ array mlx_load_npy_helper(py::object file, StreamOrDevice s) {
 }
 
 LoadOutputTypes mlx_load_helper(
-    py::object file,
+    nb::object file,
     std::optional<std::string> format,
     bool return_metadata,
     StreamOrDevice s) {
   if (!format.has_value()) {
     std::string fname;
-    if (py::isinstance<py::str>(file)) {
-      fname = py::cast<std::string>(file);
+    if (nb::isinstance<nb::str>(file)) {
+      fname = nb::cast<std::string>(file);
     } else if (is_istream_object(file)) {
-      fname = file.attr("name").cast<std::string>();
+      fname = nb::cast<std::string>(file.attr("name"));
     } else {
       throw std::invalid_argument(
           "[load] Input must be a file-like object opened in binary mode, or string");
@@ -304,14 +301,14 @@ LoadOutputTypes mlx_load_helper(
 
 class PyFileWriter : public io::Writer {
  public:
-  PyFileWriter(py::object file)
+  PyFileWriter(nb::object file)
       : pyostream_(file),
         write_func_(file.attr("write")),
         seek_func_(file.attr("seek")),
         tell_func_(file.attr("tell")) {}
 
   ~PyFileWriter() {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
 
     pyostream_.release().dec_ref();
     write_func_.release().dec_ref();
@@ -322,8 +319,8 @@ class PyFileWriter : public io::Writer {
   bool is_open() const override {
     bool out;
     {
-      py::gil_scoped_acquire gil;
-      out = !pyostream_.attr("closed").cast<bool>();
+      nb::gil_scoped_acquire gil;
+      out = !nb::cast<bool>(pyostream_.attr("closed"));
     }
     return out;
   }
@@ -331,7 +328,7 @@ class PyFileWriter : public io::Writer {
   bool good() const override {
     bool out;
     {
-      py::gil_scoped_acquire gil;
+      nb::gil_scoped_acquire gil;
       out = !pyostream_.is_none();
     }
     return out;
@@ -340,25 +337,26 @@ class PyFileWriter : public io::Writer {
   size_t tell() const override {
     size_t out;
     {
-      py::gil_scoped_acquire gil;
-      out = tell_func_().cast<size_t>();
+      nb::gil_scoped_acquire gil;
+      out = nb::cast<size_t>(tell_func_());
     }
     return out;
   }
 
   void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
       override {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
     seek_func_(off, (int)way);
   }
 
   void write(const char* data, size_t n) override {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
 
-    py::object bytes_written =
-        write_func_(py::memoryview::from_buffer(data, {n}, {sizeof(char)}));
+    auto memview =
+        PyMemoryView_FromMemory(const_cast<char*>(data), n, PyBUF_READ);
+    nb::object bytes_written = write_func_(nb::handle(memview));
 
-    if (bytes_written.is_none() || py::cast<size_t>(bytes_written) < n) {
+    if (bytes_written.is_none() || nb::cast<size_t>(bytes_written) < n) {
       throw std::runtime_error("[load] Failed to write to python stream");
     }
   }
@@ -368,20 +366,20 @@ class PyFileWriter : public io::Writer {
   }
 
  private:
-  py::object pyostream_;
-  py::object write_func_;
-  py::object seek_func_;
-  py::object tell_func_;
+  nb::object pyostream_;
+  nb::object write_func_;
+  nb::object seek_func_;
+  nb::object tell_func_;
 };
 
-void mlx_save_helper(py::object file, array a) {
-  if (py::isinstance<py::str>(file)) {
-    save(py::cast<std::string>(file), a);
+void mlx_save_helper(nb::object file, array a) {
+  if (nb::isinstance<nb::str>(file)) {
+    save(nb::cast<std::string>(file), a);
     return;
   } else if (is_ostream_object(file)) {
     auto writer = std::make_shared<PyFileWriter>(file);
     {
-      py::gil_scoped_release gil;
+      nb::gil_scoped_release gil;
       save(writer, a);
     }
 
@@ -393,26 +391,26 @@ void mlx_save_helper(py::object file, array a) {
 }
 
 void mlx_savez_helper(
-    py::object file_,
-    py::args args,
-    const py::kwargs& kwargs,
+    nb::object file_,
+    nb::args args,
+    const nb::kwargs& kwargs,
     bool compressed) {
   // Add .npz to the end of the filename if not already there
-  py::object file = file_;
+  nb::object file = file_;
 
-  if (py::isinstance<py::str>(file_)) {
-    std::string fname = file_.cast<std::string>();
+  if (nb::isinstance<nb::str>(file_)) {
+    std::string fname = nb::cast<std::string>(file_);
 
     // Add .npz to file name if it is not there
     if (fname.length() < 4 || fname.substr(fname.length() - 4, 4) != ".npz")
       fname += ".npz";
 
-    file = py::str(fname);
+    file = nb::cast(fname);
   }
 
   // Collect args and kwargs
-  auto arrays_dict = kwargs.cast<std::unordered_map<std::string, array>>();
-  auto arrays_list = args.cast<std::vector<array>>();
+  auto arrays_dict = nb::cast<std::unordered_map<std::string, array>>(kwargs);
+  auto arrays_list = nb::cast<std::vector<array>>(args);
 
   for (int i = 0; i < arrays_list.size(); i++) {
     std::string arr_name = "arr_" + std::to_string(i);
@@ -426,9 +424,9 @@ void mlx_savez_helper(
   }
 
   // Create python ZipFile object depending on compression
-  py::module_ zipfile = py::module_::import("zipfile");
-  int compression = compressed ? zipfile.attr("ZIP_DEFLATED").cast<int>()
-                               : zipfile.attr("ZIP_STORED").cast<int>();
+  nb::module_ zipfile = nb::module_::import_("zipfile");
+  int compression = nb::cast<int>(
+      compressed ? zipfile.attr("ZIP_DEFLATED") : zipfile.attr("ZIP_STORED"));
   char mode = 'w';
   ZipFileWrapper zipfile_object(zipfile, file, mode, compression);
 
@@ -438,7 +436,7 @@ void mlx_savez_helper(
     auto py_ostream = zipfile_object.open(fname, 'w');
     auto writer = std::make_shared<PyFileWriter>(py_ostream);
     {
-      py::gil_scoped_release nogil;
+      nb::gil_scoped_release nogil;
       save(writer, a);
     }
   }
@@ -447,31 +445,31 @@ void mlx_savez_helper(
 }
 
 void mlx_save_safetensor_helper(
-    py::object file,
-    py::dict d,
-    std::optional<py::dict> m) {
+    nb::object file,
+    nb::dict d,
+    std::optional<nb::dict> m) {
   std::unordered_map<std::string, std::string> metadata_map;
   if (m) {
     try {
       metadata_map =
-          m.value().cast<std::unordered_map<std::string, std::string>>();
-    } catch (const py::cast_error& e) {
+          nb::cast<std::unordered_map<std::string, std::string>>(m.value());
+    } catch (const nb::cast_error& e) {
       throw std::invalid_argument(
           "[save_safetensors] Metadata must be a dictionary with string keys and values");
     }
   } else {
     metadata_map = std::unordered_map<std::string, std::string>();
   }
-  auto arrays_map = d.cast<std::unordered_map<std::string, array>>();
-  if (py::isinstance<py::str>(file)) {
+  auto arrays_map = nb::cast<std::unordered_map<std::string, array>>(d);
+  if (nb::isinstance<nb::str>(file)) {
     {
-      py::gil_scoped_release nogil;
-      save_safetensors(py::cast<std::string>(file), arrays_map, metadata_map);
+      nb::gil_scoped_release nogil;
+      save_safetensors(nb::cast<std::string>(file), arrays_map, metadata_map);
     }
   } else if (is_ostream_object(file)) {
     auto writer = std::make_shared<PyFileWriter>(file);
     {
-      py::gil_scoped_release nogil;
+      nb::gil_scoped_release nogil;
       save_safetensors(writer, arrays_map, metadata_map);
     }
   } else {
@@ -481,22 +479,22 @@ void mlx_save_safetensor_helper(
 }
 
 void mlx_save_gguf_helper(
-    py::object file,
-    py::dict a,
-    std::optional<py::dict> m) {
-  auto arrays_map = a.cast<std::unordered_map<std::string, array>>();
-  if (py::isinstance<py::str>(file)) {
+    nb::object file,
+    nb::dict a,
+    std::optional<nb::dict> m) {
+  auto arrays_map = nb::cast<std::unordered_map<std::string, array>>(a);
+  if (nb::isinstance<nb::str>(file)) {
     if (m) {
       auto metadata_map =
-          m.value().cast<std::unordered_map<std::string, GGUFMetaData>>();
+          nb::cast<std::unordered_map<std::string, GGUFMetaData>>(m.value());
       {
-        py::gil_scoped_release nogil;
-        save_gguf(py::cast<std::string>(file), arrays_map, metadata_map);
+        nb::gil_scoped_release nogil;
+        save_gguf(nb::cast<std::string>(file), arrays_map, metadata_map);
       }
     } else {
       {
-        py::gil_scoped_release nogil;
-        save_gguf(py::cast<std::string>(file), arrays_map);
+        nb::gil_scoped_release nogil;
+        save_gguf(nb::cast<std::string>(file), arrays_map);
       }
     }
   } else {

python/src/load.h

@@ -1,15 +1,20 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/variant.h>
+
 #include <optional>
 #include <string>
 #include <unordered_map>
 #include <variant>
 #include "mlx/io.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 using namespace mlx::core;
 
 using LoadOutputTypes = std::variant<
@@ -18,27 +23,27 @@ using LoadOutputTypes = std::variant<
     SafetensorsLoad,
     GGUFLoad>;
 
-SafetensorsLoad mlx_load_safetensor_helper(py::object file, StreamOrDevice s);
+SafetensorsLoad mlx_load_safetensor_helper(nb::object file, StreamOrDevice s);
 void mlx_save_safetensor_helper(
-    py::object file,
-    py::dict d,
-    std::optional<py::dict> m);
+    nb::object file,
+    nb::dict d,
+    std::optional<nb::dict> m);
 
-GGUFLoad mlx_load_gguf_helper(py::object file, StreamOrDevice s);
+GGUFLoad mlx_load_gguf_helper(nb::object file, StreamOrDevice s);
 
 void mlx_save_gguf_helper(
-    py::object file,
-    py::dict d,
-    std::optional<py::dict> m);
+    nb::object file,
+    nb::dict d,
+    std::optional<nb::dict> m);
 
 LoadOutputTypes mlx_load_helper(
-    py::object file,
+    nb::object file,
     std::optional<std::string> format,
     bool return_metadata,
     StreamOrDevice s);
-void mlx_save_helper(py::object file, array a);
+void mlx_save_helper(nb::object file, array a);
 void mlx_savez_helper(
-    py::object file,
-    py::args args,
-    const py::kwargs& kwargs,
+    nb::object file,
+    nb::args args,
+    const nb::kwargs& kwargs,
     bool compressed = false);

python/src/metal.cpp

@@ -1,16 +1,15 @@
-// Copyright © 2023 Apple Inc.
-
-#include <pybind11/pybind11.h>
+// Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/metal/metal.h"
+#include <nanobind/nanobind.h>
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 
 using namespace mlx::core;
 
-void init_metal(py::module_& m) {
-  py::module_ metal = m.def_submodule("metal", "mlx.metal");
+void init_metal(nb::module_& m) {
+  nb::module_ metal = m.def_submodule("metal", "mlx.metal");
   metal.def(
       "is_available",
       &metal::is_available,
@@ -48,7 +47,7 @@ void init_metal(py::module_& m) {
       "set_memory_limit",
       &metal::set_memory_limit,
       "limit"_a,
-      py::kw_only(),
+      nb::kw_only(),
       "relaxed"_a = true,
       R"pbdoc(
       Set the memory limit.

python/src/mlx.cpp

@@ -1,30 +1,30 @@
-// Copyright © 2023 Apple Inc.
+// Conbright © 2023-2024 Apple Inc.
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
 
 #define STRINGIFY(x) #x
 #define TOSTRING(x) STRINGIFY(x)
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
-void init_array(py::module_&);
-void init_device(py::module_&);
-void init_stream(py::module_&);
-void init_metal(py::module_&);
-void init_ops(py::module_&);
-void init_transforms(py::module_&);
-void init_random(py::module_&);
-void init_fft(py::module_&);
-void init_linalg(py::module_&);
-void init_constants(py::module_&);
-void init_extensions(py::module_&);
-void init_utils(py::module_&);
+void init_array(nb::module_&);
+void init_device(nb::module_&);
+void init_stream(nb::module_&);
+void init_metal(nb::module_&);
+void init_ops(nb::module_&);
+void init_transforms(nb::module_&);
+void init_random(nb::module_&);
+void init_fft(nb::module_&);
+void init_linalg(nb::module_&);
+void init_constants(nb::module_&);
+void init_fast(nb::module_&);
 
-PYBIND11_MODULE(core, m) {
+NB_MODULE(core, m) {
   m.doc() = "mlx: A framework for machine learning on Apple silicon.";
 
-  auto reprlib_fix = py::module_::import("mlx._reprlib_fix");
-  py::module_::import("mlx._os_warning");
+  auto reprlib_fix = nb::module_::import_("mlx._reprlib_fix");
+  nb::module_::import_("mlx._os_warning");
+  nb::set_leak_warnings(false);
 
   init_device(m);
   init_stream(m);
@@ -36,8 +36,7 @@ PYBIND11_MODULE(core, m) {
   init_fft(m);
   init_linalg(m);
   init_constants(m);
-  init_extensions(m);
-  init_utils(m);
+  init_fast(m);
 
   m.attr("__version__") = TOSTRING(_VERSION_);
 }

python/src/pybind11_numpy_fp16.h

@@ -1,60 +0,0 @@
-// Copyright © 2023-2024 Apple Inc.
-#pragma once
-
-// A patch to get float16_t to work with pybind11 numpy arrays
-// Derived from:
-// https://github.com/pybind/pybind11/issues/1776#issuecomment-492230679
-
-#include <pybind11/numpy.h>
-
-namespace pybind11::detail {
-
-template <typename T>
-struct npy_scalar_caster {
-  PYBIND11_TYPE_CASTER(T, _("PleaseOverride"));
-  using Array = array_t<T>;
-
-  bool load(handle src, bool convert) {
-    // Taken from Eigen casters. Permits either scalar dtype or scalar array.
-    handle type = dtype::of<T>().attr("type"); // Could make more efficient.
-    if (!convert && !isinstance<Array>(src) && !isinstance(src, type))
-      return false;
-    Array tmp = Array::ensure(src);
-    if (tmp && tmp.size() == 1 && tmp.ndim() == 0) {
-      this->value = *tmp.data();
-      return true;
-    }
-    return false;
-  }
-
-  static handle cast(T src, return_value_policy, handle) {
-    Array tmp({1});
-    tmp.mutable_at(0) = src;
-    tmp.resize({});
-    // You could also just return the array if you want a scalar array.
-    object scalar = tmp[tuple()];
-    return scalar.release();
-  }
-};
-
-// Similar to enums in `pybind11/numpy.h`. Determined by doing:
-// python3 -c 'import numpy as np; print(np.dtype(np.float16).num)'
-constexpr int NPY_FLOAT16 = 23;
-
-// Kinda following:
-// https://github.com/pybind/pybind11/blob/9bb3313162c0b856125e481ceece9d8faa567716/include/pybind11/numpy.h#L1000
-template <>
-struct npy_format_descriptor<float16_t> {
-  static constexpr auto name = _("float16");
-  static pybind11::dtype dtype() {
-    handle ptr = npy_api::get().PyArray_DescrFromType_(NPY_FLOAT16);
-    return reinterpret_borrow<pybind11::dtype>(ptr);
-  }
-};
-
-template <>
-struct type_caster<float16_t> : npy_scalar_caster<float16_t> {
-  static constexpr auto name = _("float16");
-};
-
-} // namespace pybind11::detail

python/src/random.cpp

@@ -1,7 +1,10 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/variant.h>
+#include <nanobind/stl/vector.h>
 
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
 #include <chrono>
 
 #include "python/src/utils.h"
@@ -9,8 +12,8 @@
 #include "mlx/ops.h"
 #include "mlx/random.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 using namespace mlx::core::random;
 
@@ -25,22 +28,22 @@ class PyKeySequence {
   }
 
   array next() {
-    auto out = split(py::cast<array>(state_[0]));
+    auto out = split(nb::cast<array>(state_[0]));
     state_[0] = out.first;
     return out.second;
   }
 
-  py::list state() {
+  nb::list state() {
     return state_;
   }
 
   void release() {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
     state_.release().dec_ref();
   }
 
  private:
-  py::list state_;
+  nb::list state_;
 };
 
 PyKeySequence& default_key() {
@@ -54,7 +57,7 @@ PyKeySequence& default_key() {
   return ks;
 }
 
-void init_random(py::module_& parent_module) {
+void init_random(nb::module_& parent_module) {
   auto m = parent_module.def_submodule(
       "random",
       "mlx.core.random: functionality related to random number generation");
@@ -85,10 +88,10 @@ void init_random(py::module_& parent_module) {
       )pbdoc");
   m.def(
       "split",
-      py::overload_cast<const array&, int, StreamOrDevice>(&random::split),
+      nb::overload_cast<const array&, int, StreamOrDevice>(&random::split),
       "key"_a,
       "num"_a = 2,
-      "stream"_a = none,
+      "stream"_a = nb::none(),
       R"pbdoc(
         Split a PRNG key into sub keys.
 
@@ -119,9 +122,9 @@ void init_random(py::module_& parent_module) {
       "low"_a = 0,
       "high"_a = 1,
       "shape"_a = std::vector<int>{},
-      "dtype"_a = std::optional{float32},
-      "key"_a = none,
-      "stream"_a = none,
+      "dtype"_a.none() = float32,
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Generate uniformly distributed random numbers.
 
@@ -151,11 +154,11 @@ void init_random(py::module_& parent_module) {
         return normal(shape, type.value_or(float32), loc, scale, key, s);
       },
       "shape"_a = std::vector<int>{},
-      "dtype"_a = std::optional{float32},
+      "dtype"_a.none() = float32,
       "loc"_a = 0.0,
       "scale"_a = 1.0,
-      "key"_a = none,
-      "stream"_a = none,
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Generate normally distributed random numbers.
 
@@ -184,9 +187,9 @@ void init_random(py::module_& parent_module) {
       "low"_a,
       "high"_a,
       "shape"_a = std::vector<int>{},
-      "dtype"_a = int32,
-      "key"_a = none,
-      "stream"_a = none,
+      "dtype"_a.none() = int32,
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Generate random integers from the given interval.
 
@@ -219,9 +222,9 @@ void init_random(py::module_& parent_module) {
         }
       },
       "p"_a = 0.5,
-      "shape"_a = none,
-      "key"_a = none,
-      "stream"_a = none,
+      "shape"_a = nb::none(),
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Generate Bernoulli random values.
 
@@ -259,10 +262,10 @@ void init_random(py::module_& parent_module) {
       },
       "lower"_a,
       "upper"_a,
-      "shape"_a = none,
-      "dtype"_a = std::optional{float32},
-      "key"_a = none,
-      "stream"_a = none,
+      "shape"_a = nb::none(),
+      "dtype"_a.none() = float32,
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Generate values from a truncated normal distribution.
 
@@ -292,9 +295,9 @@ void init_random(py::module_& parent_module) {
         return gumbel(shape, type.value_or(float32), key, s);
       },
       "shape"_a = std::vector<int>{},
-      "dtype"_a = std::optional{float32},
-      "stream"_a = none,
-      "key"_a = none,
+      "dtype"_a.none() = float32,
+      "stream"_a = nb::none(),
+      "key"_a = nb::none(),
       R"pbdoc(
         Sample from the standard Gumbel distribution.
 
@@ -331,10 +334,10 @@ void init_random(py::module_& parent_module) {
       },
       "logits"_a,
       "axis"_a = -1,
-      "shape"_a = none,
-      "num_samples"_a = none,
-      "key"_a = none,
-      "stream"_a = none,
+      "shape"_a = nb::none(),
+      "num_samples"_a = nb::none(),
+      "key"_a = nb::none(),
+      "stream"_a = nb::none(),
       R"pbdoc(
         Sample from a categorical distribution.
 
@@ -359,6 +362,6 @@ void init_random(py::module_& parent_module) {
             array: The ``shape``-sized output array with type ``uint32``.
       )pbdoc");
   // Register static Python object cleanup before the interpreter exits
-  auto atexit = py::module_::import("atexit");
-  atexit.attr("register")(py::cpp_function([]() { default_key().release(); }));
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function([]() { default_key().release(); }));
 }

python/src/stream.cpp

@@ -1,25 +1,54 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <sstream>
 
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/stream.h"
 #include "mlx/utils.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 
-void init_stream(py::module_& m) {
-  py::class_<Stream>(
+// Create the StreamContext on enter and delete on exit.
+class PyStreamContext {
+ public:
+  PyStreamContext(StreamOrDevice s) : _inner(nullptr) {
+    if (std::holds_alternative<std::monostate>(s)) {
+      throw std::runtime_error(
+          "[StreamContext] Invalid argument, please specify a stream or device.");
+    }
+    _s = s;
+  }
+
+  void enter() {
+    _inner = new StreamContext(_s);
+  }
+
+  void exit() {
+    if (_inner != nullptr) {
+      delete _inner;
+      _inner = nullptr;
+    }
+  }
+
+ private:
+  StreamOrDevice _s;
+  StreamContext* _inner;
+};
+
+void init_stream(nb::module_& m) {
+  nb::class_<Stream>(
       m,
       "Stream",
       R"pbdoc(
       A stream for running operations on a given device.
       )pbdoc")
-      .def(py::init<int, Device>(), "index"_a, "device"_a)
-      .def_readonly("device", &Stream::device)
+      .def(nb::init<int, Device>(), "index"_a, "device"_a)
+      .def_ro("device", &Stream::device)
       .def(
           "__repr__",
           [](const Stream& s) {
@@ -31,7 +60,7 @@ void init_stream(py::module_& m) {
         return s1 == s2;
       });
 
-  py::implicitly_convertible<Device::DeviceType, Device>();
+  nb::implicitly_convertible<Device::DeviceType, Device>();
 
   m.def(
       "default_stream",
@@ -56,4 +85,48 @@ void init_stream(py::module_& m) {
       &new_stream,
       "device"_a,
       R"pbdoc(Make a new stream on the given device.)pbdoc");
+
+  nb::class_<PyStreamContext>(m, "StreamContext", R"pbdoc(
+        A context manager for setting the current device and stream.
+
+        See :func:`stream` for usage.
+
+        Args:
+            s: The stream or device to set as the default.
+  )pbdoc")
+      .def(nb::init<StreamOrDevice>(), "s"_a)
+      .def("__enter__", [](PyStreamContext& scm) { scm.enter(); })
+      .def(
+          "__exit__",
+          [](PyStreamContext& scm,
+             const std::optional<nb::type_object>& exc_type,
+             const std::optional<nb::object>& exc_value,
+             const std::optional<nb::object>& traceback) { scm.exit(); },
+          "exc_type"_a = nb::none(),
+          "exc_value"_a = nb::none(),
+          "traceback"_a = nb::none());
+  m.def(
+      "stream",
+      [](StreamOrDevice s) { return PyStreamContext(s); },
+      "s"_a,
+      R"pbdoc(
+        Create a context manager to set the default device and stream.
+
+        Args:
+            s: The :obj:`Stream` or :obj:`Device` to set as the default.
+
+        Returns:
+            A context manager that sets the default device and stream.
+
+        Example:
+
+        .. code-block::python
+
+          import mlx.core as mx
+
+          # Create a context manager for the default device and stream.
+          with mx.stream(mx.cpu):
+              # Operations here will use mx.cpu by default.
+              pass
+      )pbdoc");
 }

python/src/transforms.cpp

@@ -1,6 +1,11 @@
 // Copyright © 2023-2024 Apple Inc.
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/pair.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/variant.h>
+#include <nanobind/stl/vector.h>
+
 #include <algorithm>
 #include <fstream>
 #include <numeric>
@@ -13,13 +18,17 @@
 #include "mlx/transforms_impl.h"
 #include "python/src/trees.h"
 
-namespace py = pybind11;
-using namespace py::literals;
+namespace nb = nanobind;
+using namespace nb::literals;
 using namespace mlx::core;
 
 using IntOrVec = std::variant<int, std::vector<int>>;
 using StrOrVec = std::variant<std::string, std::vector<std::string>>;
 
+inline std::string type_name_str(const nb::handle& o) {
+  return nb::cast<std::string>(nb::type_name(o.type()));
+}
+
 template <typename T>
 std::vector<T> to_vector(const std::variant<T, std::vector<T>>& v) {
   std::vector<T> vals;
@@ -49,7 +58,7 @@ auto validate_argnums_argnames(
 }
 
 auto py_value_and_grad(
-    const py::function& fun,
+    const nb::callable& fun,
     std::vector<int> argnums,
     std::vector<std::string> argnames,
     const std::string& error_msg_tag,
@@ -71,7 +80,7 @@ auto py_value_and_grad(
   }
 
   return [fun, argnums, argnames, error_msg_tag, scalar_func_only](
-             const py::args& args, const py::kwargs& kwargs) {
+             const nb::args& args, const nb::kwargs& kwargs) {
     // Sanitize the input
     if (argnums.size() > 0 && argnums.back() >= args.size()) {
       std::ostringstream msg;
@@ -89,7 +98,7 @@ auto py_value_and_grad(
             << "' because the function is called with the "
             << "following keyword arguments {";
         for (auto item : kwargs) {
-          msg << item.first.cast<std::string>() << ",";
+          msg << nb::cast<std::string>(item.first) << ",";
         }
         msg << "}";
         throw std::invalid_argument(msg.str());
@@ -115,7 +124,7 @@ auto py_value_and_grad(
 
     // value_out will hold the output of the python function in order to be
     // able to reconstruct the python tree of extra return values
-    py::object py_value_out;
+    nb::object py_value_out;
     auto value_and_grads = value_and_grad(
         [&fun,
          &args,
@@ -127,15 +136,15 @@ auto py_value_and_grad(
          &error_msg_tag,
          scalar_func_only](const std::vector<array>& a) {
           // Copy the arguments
-          py::args args_cpy = py::tuple(args.size());
-          py::kwargs kwargs_cpy = py::kwargs();
+          nb::list args_cpy;
+          nb::kwargs kwargs_cpy = nb::kwargs();
           int j = 0;
           for (int i = 0; i < args.size(); ++i) {
             if (j < argnums.size() && i == argnums[j]) {
-              args_cpy[i] = tree_unflatten(args[i], a, counts[j]);
+              args_cpy.append(tree_unflatten(args[i], a, counts[j]));
               j++;
             } else {
-              args_cpy[i] = args[i];
+              args_cpy.append(args[i]);
             }
           }
           for (auto& key : argnames) {
@@ -154,25 +163,25 @@ auto py_value_and_grad(
           py_value_out = fun(*args_cpy, **kwargs_cpy);
 
           // Validate the return value of the python function
-          if (!py::isinstance<array>(py_value_out)) {
+          if (!nb::isinstance<array>(py_value_out)) {
             if (scalar_func_only) {
               std::ostringstream msg;
               msg << error_msg_tag << " The return value of the function "
                   << "whose gradient we want to compute should be a "
-                  << "scalar array; but " << py_value_out.get_type()
+                  << "scalar array; but " << type_name_str(py_value_out)
                   << " was returned.";
               throw std::invalid_argument(msg.str());
             }
-            if (!py::isinstance<py::tuple>(py_value_out)) {
+            if (!nb::isinstance<nb::tuple>(py_value_out)) {
               std::ostringstream msg;
               msg << error_msg_tag << " The return value of the function "
                   << "whose gradient we want to compute should be either a "
                   << "scalar array or a tuple with the first value being a "
                   << "scalar array (Union[array, Tuple[array, Any, ...]]); but "
-                  << py_value_out.get_type() << " was returned.";
+                  << type_name_str(py_value_out) << " was returned.";
               throw std::invalid_argument(msg.str());
             }
-            py::tuple ret = py::cast<py::tuple>(py_value_out);
+            nb::tuple ret = nb::cast<nb::tuple>(py_value_out);
             if (ret.size() == 0) {
               std::ostringstream msg;
               msg << error_msg_tag << " The return value of the function "
@@ -182,14 +191,14 @@ auto py_value_and_grad(
                   << "we got an empty tuple.";
               throw std::invalid_argument(msg.str());
             }
-            if (!py::isinstance<array>(ret[0])) {
+            if (!nb::isinstance<array>(ret[0])) {
               std::ostringstream msg;
               msg << error_msg_tag << " The return value of the function "
                   << "whose gradient we want to compute should be either a "
                   << "scalar array or a tuple with the first value being a "
                   << "scalar array (Union[array, Tuple[array, Any, ...]]); but it "
                   << "was a tuple with the first value being of type "
-                  << ret[0].get_type() << " .";
+                  << type_name_str(ret[0]) << " .";
               throw std::invalid_argument(msg.str());
             }
           }
@@ -212,61 +221,60 @@ auto py_value_and_grad(
     // In case 2 we return a tuple of the above.
     // In case 3 we return a tuple containing a tuple and dict (sth like
     // (tuple(), dict(x=mx.array(5))) ).
-    py::object positional_grads;
-    py::object keyword_grads;
-    py::object py_grads;
+    nb::object positional_grads;
+    nb::object keyword_grads;
+    nb::object py_grads;
 
     // Collect the gradients for the positional arguments
     if (argnums.size() == 1) {
       positional_grads = tree_unflatten(args[argnums[0]], gradients, counts[0]);
     } else if (argnums.size() > 1) {
-      py::tuple grads_(argnums.size());
+      nb::list grads_;
       for (int i = 0; i < argnums.size(); i++) {
-        grads_[i] = tree_unflatten(args[argnums[i]], gradients, counts[i]);
+        grads_.append(tree_unflatten(args[argnums[i]], gradients, counts[i]));
       }
-      positional_grads = py::cast<py::object>(grads_);
+      positional_grads = nb::tuple(grads_);
     } else {
-      positional_grads = py::none();
+      positional_grads = nb::none();
     }
 
     // No keyword argument gradients so return the tuple of gradients
     if (argnames.size() == 0) {
       py_grads = positional_grads;
     } else {
-      py::dict grads_;
+      nb::dict grads_;
       for (int i = 0; i < argnames.size(); i++) {
         auto& k = argnames[i];
         grads_[k.c_str()] = tree_unflatten(
             kwargs[k.c_str()], gradients, counts[i + argnums.size()]);
       }
-      keyword_grads = py::cast<py::object>(grads_);
+      keyword_grads = grads_;
 
-      py_grads =
-          py::cast<py::object>(py::make_tuple(positional_grads, keyword_grads));
+      py_grads = nb::make_tuple(positional_grads, keyword_grads);
     }
 
     // Put the values back in the container
-    py::object return_value = tree_unflatten(py_value_out, value);
+    nb::object return_value = tree_unflatten(py_value_out, value);
     return std::make_pair(return_value, py_grads);
   };
 }
 
 auto py_vmap(
-    const py::function& fun,
-    const py::object& in_axes,
-    const py::object& out_axes) {
-  return [fun, in_axes, out_axes](const py::args& args) {
-    auto axes_to_flat_tree = [](const py::object& tree,
-                                const py::object& axes) {
+    const nb::callable& fun,
+    const nb::object& in_axes,
+    const nb::object& out_axes) {
+  return [fun, in_axes, out_axes](const nb::args& args) {
+    auto axes_to_flat_tree = [](const nb::object& tree,
+                                const nb::object& axes) {
       auto tree_axes = tree_map(
           {tree, axes},
-          [](const std::vector<py::object>& inputs) { return inputs[1]; });
+          [](const std::vector<nb::object>& inputs) { return inputs[1]; });
       std::vector<int> flat_axes;
-      tree_visit(tree_axes, [&flat_axes](py::handle obj) {
+      tree_visit(tree_axes, [&flat_axes](nb::handle obj) {
         if (obj.is_none()) {
           flat_axes.push_back(-1);
-        } else if (py::isinstance<py::int_>(obj)) {
-          flat_axes.push_back(py::cast<int>(py::cast<py::int_>(obj)));
+        } else if (nb::isinstance<nb::int_>(obj)) {
+          flat_axes.push_back(nb::cast<int>(nb::cast<nb::int_>(obj)));
         } else {
           throw std::invalid_argument("[vmap] axis must be int or None.");
         }
@@ -280,7 +288,7 @@ auto py_vmap(
 
     // py_value_out will hold the output of the python function in order to be
     // able to reconstruct the python tree of extra return values
-    py::object py_outputs;
+    nb::object py_outputs;
 
     auto vmap_fn =
         [&fun, &args, &inputs, &py_outputs](const std::vector<array>& a) {
@@ -305,24 +313,24 @@ auto py_vmap(
   };
 }
 
-std::unordered_map<size_t, py::object>& tree_cache() {
+std::unordered_map<size_t, nb::object>& tree_cache() {
   // This map is used to Cache the tree structure of the outputs
-  static std::unordered_map<size_t, py::object> tree_cache_;
+  static std::unordered_map<size_t, nb::object> tree_cache_;
   return tree_cache_;
 }
 
 struct PyCompiledFun {
-  py::function fun;
+  nb::callable fun;
   size_t fun_id;
-  py::object captured_inputs;
-  py::object captured_outputs;
+  nb::object captured_inputs;
+  nb::object captured_outputs;
   bool shapeless;
-  size_t num_outputs{0};
+  mutable size_t num_outputs{0};
 
   PyCompiledFun(
-      const py::function& fun,
-      py::object inputs,
-      py::object outputs,
+      const nb::callable& fun,
+      nb::object inputs,
+      nb::object outputs,
       bool shapeless)
       : fun(fun),
         fun_id(reinterpret_cast<size_t>(fun.ptr())),
@@ -342,7 +350,7 @@ struct PyCompiledFun {
     num_outputs = other.num_outputs;
   };
 
-  py::object operator()(const py::args& args, const py::kwargs& kwargs) {
+  nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) {
     // Flat array inputs
     std::vector<array> inputs;
 
@@ -358,45 +366,45 @@ struct PyCompiledFun {
     constexpr uint64_t dict_identifier = 18446744073709551521UL;
 
     // Flatten the tree with hashed constants and structure
-    std::function<void(py::handle)> recurse;
-    recurse = [&](py::handle obj) {
-      if (py::isinstance<py::list>(obj)) {
-        auto l = py::cast<py::list>(obj);
+    std::function<void(nb::handle)> recurse;
+    recurse = [&](nb::handle obj) {
+      if (nb::isinstance<nb::list>(obj)) {
+        auto l = nb::cast<nb::list>(obj);
         constants.push_back(list_identifier);
         for (int i = 0; i < l.size(); ++i) {
           recurse(l[i]);
         }
-      } else if (py::isinstance<py::tuple>(obj)) {
-        auto l = py::cast<py::tuple>(obj);
+      } else if (nb::isinstance<nb::tuple>(obj)) {
+        auto l = nb::cast<nb::tuple>(obj);
         constants.push_back(list_identifier);
         for (auto item : obj) {
           recurse(item);
         }
-      } else if (py::isinstance<py::dict>(obj)) {
-        auto d = py::cast<py::dict>(obj);
+      } else if (nb::isinstance<nb::dict>(obj)) {
+        auto d = nb::cast<nb::dict>(obj);
         constants.push_back(dict_identifier);
         for (auto item : d) {
-          auto r = py::hash(item.first);
+          auto r = item.first.attr("__hash__");
           constants.push_back(*reinterpret_cast<uint64_t*>(&r));
           recurse(item.second);
         }
-      } else if (py::isinstance<array>(obj)) {
-        inputs.push_back(py::cast<array>(obj));
+      } else if (nb::isinstance<array>(obj)) {
+        inputs.push_back(nb::cast<array>(obj));
         constants.push_back(array_identifier);
-      } else if (py::isinstance<py::str>(obj)) {
-        auto r = py::hash(obj);
+      } else if (nb::isinstance<nb::str>(obj)) {
+        auto r = obj.attr("__hash__");
         constants.push_back(*reinterpret_cast<uint64_t*>(&r));
-      } else if (py::isinstance<py::int_>(obj)) {
-        auto r = obj.cast<int64_t>();
+      } else if (nb::isinstance<nb::int_>(obj)) {
+        auto r = nb::cast<int64_t>(obj);
         constants.push_back(*reinterpret_cast<uint64_t*>(&r));
-      } else if (py::isinstance<py::float_>(obj)) {
-        auto r = obj.cast<double>();
+      } else if (nb::isinstance<nb::float_>(obj)) {
+        auto r = nb::cast<double>(obj);
         constants.push_back(*reinterpret_cast<uint64_t*>(&r));
       } else {
         std::ostringstream msg;
         msg << "[compile] Function arguments must be trees of arrays "
             << "or constants (floats, ints, or strings), but received "
-            << "type " << obj.get_type() << ".";
+            << "type " << type_name_str(obj) << ".";
         throw std::invalid_argument(msg.str());
       }
     };
@@ -404,13 +412,12 @@ struct PyCompiledFun {
     recurse(args);
     int num_args = inputs.size();
     recurse(kwargs);
-
     auto compile_fun = [this, &args, &kwargs, num_args](
                            const std::vector<array>& a) {
       // Put tracers into captured inputs
       std::vector<array> flat_in_captures;
       std::vector<array> trace_captures;
-      if (!py::isinstance<py::none>(captured_inputs)) {
+      if (!captured_inputs.is_none()) {
         flat_in_captures = tree_flatten(captured_inputs, false);
         trace_captures.insert(
             trace_captures.end(), a.end() - flat_in_captures.size(), a.end());
@@ -425,7 +432,7 @@ struct PyCompiledFun {
       tree_cache().insert({fun_id, py_outputs});
 
       num_outputs = outputs.size();
-      if (!py::isinstance<py::none>(captured_outputs)) {
+      if (!captured_outputs.is_none()) {
         auto flat_out_captures = tree_flatten(captured_outputs, false);
         outputs.insert(
             outputs.end(),
@@ -434,13 +441,13 @@ struct PyCompiledFun {
       }
 
       // Replace tracers with originals in captured inputs
-      if (!py::isinstance<py::none>(captured_inputs)) {
+      if (!captured_inputs.is_none()) {
         tree_replace(captured_inputs, trace_captures, flat_in_captures);
       }
       return outputs;
     };
 
-    if (!py::isinstance<py::none>(captured_inputs)) {
+    if (!captured_inputs.is_none()) {
       auto flat_in_captures = tree_flatten(captured_inputs, false);
       inputs.insert(
           inputs.end(),
@@ -451,7 +458,7 @@ struct PyCompiledFun {
     // Compile and call
     auto outputs =
         detail::compile(compile_fun, fun_id, shapeless, constants)(inputs);
-    if (!py::isinstance<py::none>(captured_outputs)) {
+    if (!captured_outputs.is_none()) {
       std::vector<array> captures(
           std::make_move_iterator(outputs.begin() + num_outputs),
           std::make_move_iterator(outputs.end()));
@@ -459,12 +466,16 @@ struct PyCompiledFun {
     }
 
     // Put the outputs back in the container
-    py::object py_outputs = tree_cache().at(fun_id);
+    nb::object py_outputs = tree_cache().at(fun_id);
     return tree_unflatten_from_structure(py_outputs, outputs);
+  }
+
+  nb::object operator()(const nb::args& args, const nb::kwargs& kwargs) const {
+    return const_cast<PyCompiledFun*>(this)->call_impl(args, kwargs);
   };
 
   ~PyCompiledFun() {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
 
     tree_cache().erase(fun_id);
     detail::compile_erase(fun_id);
@@ -476,35 +487,35 @@ struct PyCompiledFun {
 
 class PyCheckpointedFun {
  public:
-  PyCheckpointedFun(py::function fun) : fun_(std::move(fun)) {}
+  PyCheckpointedFun(nb::callable fun) : fun_(std::move(fun)) {}
 
   ~PyCheckpointedFun() {
-    py::gil_scoped_acquire gil;
+    nb::gil_scoped_acquire gil;
 
     fun_.release().dec_ref();
   }
 
   struct InnerFunction {
-    py::object fun_;
-    py::object args_structure_;
-    std::weak_ptr<py::object> output_structure_;
+    nb::object fun_;
+    nb::object args_structure_;
+    std::weak_ptr<nb::object> output_structure_;
 
     InnerFunction(
-        py::object fun,
-        py::object args_structure,
-        std::weak_ptr<py::object> output_structure)
+        nb::object fun,
+        nb::object args_structure,
+        std::weak_ptr<nb::object> output_structure)
         : fun_(std::move(fun)),
           args_structure_(std::move(args_structure)),
           output_structure_(output_structure) {}
     ~InnerFunction() {
-      py::gil_scoped_acquire gil;
+      nb::gil_scoped_acquire gil;
 
       fun_.release().dec_ref();
       args_structure_.release().dec_ref();
     }
 
     std::vector<array> operator()(const std::vector<array>& inputs) {
-      auto args = py::cast<py::tuple>(
+      auto args = nb::cast<nb::tuple>(
           tree_unflatten_from_structure(args_structure_, inputs));
       auto [outputs, output_structure] =
           tree_flatten_with_structure(fun_(*args[0], **args[1]), false);
@@ -515,9 +526,9 @@ class PyCheckpointedFun {
     }
   };
 
-  py::object operator()(const py::args& args, const py::kwargs& kwargs) {
-    auto output_structure = std::make_shared<py::object>();
-    auto full_args = py::make_tuple(args, kwargs);
+  nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) {
+    auto output_structure = std::make_shared<nb::object>();
+    auto full_args = nb::make_tuple(args, kwargs);
     auto [inputs, args_structure] =
         tree_flatten_with_structure(full_args, false);
 
@@ -527,26 +538,27 @@ class PyCheckpointedFun {
     return tree_unflatten_from_structure(*output_structure, outputs);
   }
 
+  nb::object operator()(const nb::args& args, const nb::kwargs& kwargs) const {
+    return const_cast<PyCheckpointedFun*>(this)->call_impl(args, kwargs);
+  }
+
  private:
-  py::function fun_;
+  nb::callable fun_;
 };
 
-void init_transforms(py::module_& m) {
-  py::options options;
-  options.disable_function_signatures();
-
+void init_transforms(nb::module_& m) {
   m.def(
       "eval",
-      [](const py::args& args) {
+      [](const nb::args& args) {
         std::vector<array> arrays = tree_flatten(args, false);
         {
-          py::gil_scoped_release nogil;
+          nb::gil_scoped_release nogil;
           eval(arrays);
         }
       },
+      nb::arg(),
+      nb::sig("def eval(*args) -> None"),
       R"pbdoc(
-        eval(*args) -> None
-
         Evaluate an :class:`array` or tree of :class:`array`.
 
         Args:
@@ -557,19 +569,15 @@ void init_transforms(py::module_& m) {
       )pbdoc");
   m.def(
       "jvp",
-      [](const py::function& fun,
+      [](const nb::callable& fun,
          const std::vector<array>& primals,
          const std::vector<array>& tangents) {
         auto vfun = [&fun](const std::vector<array>& primals) {
-          py::args args = py::tuple(primals.size());
-          for (int i = 0; i < primals.size(); ++i) {
-            args[i] = primals[i];
-          }
-          auto out = fun(*args);
-          if (py::isinstance<array>(out)) {
-            return std::vector<array>{py::cast<array>(out)};
+          auto out = fun(*nb::cast(primals));
+          if (nb::isinstance<array>(out)) {
+            return std::vector<array>{nb::cast<array>(out)};
           } else {
-            return py::cast<std::vector<array>>(out);
+            return nb::cast<std::vector<array>>(out);
           }
         };
         return jvp(vfun, primals, tangents);
@@ -577,17 +585,16 @@ void init_transforms(py::module_& m) {
       "fun"_a,
       "primals"_a,
       "tangents"_a,
+      nb::sig(
+          "def jvp(fun: callable, primals: List[array], tangents: List[array]) -> Tuple[List[array], List[array]]"),
       R"pbdoc(
-        jvp(fun: function, primals: List[array], tangents: List[array]) -> Tuple[List[array], List[array]]
-
-
         Compute the Jacobian-vector product.
 
         This computes the product of the Jacobian of a function ``fun`` evaluated
         at ``primals`` with the ``tangents``.
 
         Args:
-            fun (function): A function which takes a variable number of :class:`array`
+            fun (callable): A function which takes a variable number of :class:`array`
               and returns a single :class:`array` or list of :class:`array`.
             primals (list(array)): A list of :class:`array` at which to
               evaluate the Jacobian.
@@ -601,19 +608,15 @@ void init_transforms(py::module_& m) {
       )pbdoc");
   m.def(
       "vjp",
-      [](const py::function& fun,
+      [](const nb::callable& fun,
          const std::vector<array>& primals,
          const std::vector<array>& cotangents) {
         auto vfun = [&fun](const std::vector<array>& primals) {
-          py::args args = py::tuple(primals.size());
-          for (int i = 0; i < primals.size(); ++i) {
-            args[i] = primals[i];
-          }
-          auto out = fun(*args);
-          if (py::isinstance<array>(out)) {
-            return std::vector<array>{py::cast<array>(out)};
+          auto out = fun(*nb::cast(primals));
+          if (nb::isinstance<array>(out)) {
+            return std::vector<array>{nb::cast<array>(out)};
           } else {
-            return py::cast<std::vector<array>>(out);
+            return nb::cast<std::vector<array>>(out);
           }
         };
         return vjp(vfun, primals, cotangents);
@@ -621,16 +624,16 @@ void init_transforms(py::module_& m) {
       "fun"_a,
       "primals"_a,
       "cotangents"_a,
+      nb::sig(
+          "def vjp(fun: callable, primals: List[array], cotangents: List[array]) -> Tuple[List[array], List[array]]"),
       R"pbdoc(
-        vjp(fun: function, primals: List[array], cotangents: List[array]) -> Tuple[List[array], List[array]]
-
         Compute the vector-Jacobian product.
 
         Computes the product of the ``cotangents`` with the Jacobian of a
         function ``fun`` evaluated at ``primals``.
 
         Args:
-          fun (function): A function which takes a variable number of :class:`array`
+          fun (callable): A function which takes a variable number of :class:`array`
             and returns a single :class:`array` or list of :class:`array`.
           primals (list(array)): A list of :class:`array` at which to
             evaluate the Jacobian.
@@ -644,20 +647,20 @@ void init_transforms(py::module_& m) {
       )pbdoc");
   m.def(
       "value_and_grad",
-      [](const py::function& fun,
+      [](const nb::callable& fun,
          const std::optional<IntOrVec>& argnums,
          const StrOrVec& argnames) {
         auto [argnums_vec, argnames_vec] =
             validate_argnums_argnames(argnums, argnames);
-        return py::cpp_function(py_value_and_grad(
+        return nb::cpp_function(py_value_and_grad(
             fun, argnums_vec, argnames_vec, "[value_and_grad]", false));
       },
       "fun"_a,
-      "argnums"_a = std::nullopt,
+      "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
+      nb::sig(
+          "def value_and_grad(fun: callable, argnums: Optional[Union[int, List[int]]] = None, argnames: Union[str, List[str]] = []) -> callable"),
       R"pbdoc(
-        value_and_grad(fun: function, argnums: Optional[Union[int, List[int]]] = None, argnames: Union[str, List[str]] = []) -> function
-
         Returns a function which computes the value and gradient of ``fun``.
 
         The function passed to :func:`value_and_grad` should return either
@@ -688,7 +691,7 @@ void init_transforms(py::module_& m) {
             (loss, mse, l1), grads = mx.value_and_grad(lasso)(params, inputs, targets)
 
         Args:
-            fun (function): A function which takes a variable number of
+            fun (callable): A function which takes a variable number of
               :class:`array` or trees of :class:`array` and returns
               a scalar output :class:`array` or a tuple the first element
               of which should be a scalar :class:`array`.
@@ -702,34 +705,34 @@ void init_transforms(py::module_& m) {
               no gradients for keyword arguments by default.
 
         Returns:
-            function: A function which returns a tuple where the first element
+            callable: A function which returns a tuple where the first element
             is the output of `fun` and the second element is the gradients w.r.t.
             the loss.
       )pbdoc");
   m.def(
       "grad",
-      [](const py::function& fun,
+      [](const nb::callable& fun,
          const std::optional<IntOrVec>& argnums,
          const StrOrVec& argnames) {
         auto [argnums_vec, argnames_vec] =
             validate_argnums_argnames(argnums, argnames);
         auto fn =
             py_value_and_grad(fun, argnums_vec, argnames_vec, "[grad]", true);
-        return py::cpp_function(
-            [fn](const py::args& args, const py::kwargs& kwargs) {
+        return nb::cpp_function(
+            [fn](const nb::args& args, const nb::kwargs& kwargs) {
               return fn(args, kwargs).second;
             });
       },
       "fun"_a,
-      "argnums"_a = std::nullopt,
+      "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
+      nb::sig(
+          "def grad(fun: callable, argnums: Optional[Union[int, List[int]]] = None, argnames: Union[str, List[str]] = []) -> callable"),
       R"pbdoc(
-        grad(fun: function, argnums: Optional[Union[int, List[int]]] = None, argnames: Union[str, List[str]] = []) -> function
-
         Returns a function which computes the gradient of ``fun``.
 
         Args:
-            fun (function): A function which takes a variable number of
+            fun (callable): A function which takes a variable number of
               :class:`array` or trees of :class:`array` and returns
               a scalar output :class:`array`.
             argnums (int or list(int), optional): Specify the index (or indices)
@@ -742,26 +745,26 @@ void init_transforms(py::module_& m) {
               no gradients for keyword arguments by default.
 
         Returns:
-            function: A function which has the same input arguments as ``fun`` and
+            callable: A function which has the same input arguments as ``fun`` and
             returns the gradient(s).
       )pbdoc");
   m.def(
       "vmap",
-      [](const py::function& fun,
-         const py::object& in_axes,
-         const py::object& out_axes) {
-        return py::cpp_function(py_vmap(fun, in_axes, out_axes));
+      [](const nb::callable& fun,
+         const nb::object& in_axes,
+         const nb::object& out_axes) {
+        return nb::cpp_function(py_vmap(fun, in_axes, out_axes));
       },
       "fun"_a,
       "in_axes"_a = 0,
       "out_axes"_a = 0,
+      nb::sig(
+          "def vmap(fun: callable, in_axes: object = 0, out_axes: object = 0) -> callable"),
       R"pbdoc(
-        vmap(fun: function, in_axes: object = 0, out_axes: object = 0) -> function
-
         Returns a vectorized version of ``fun``.
 
         Args:
-            fun (function): A function which takes a variable number of
+            fun (callable): A function which takes a variable number of
               :class:`array` or a tree of :class:`array` and returns
               a variable number of :class:`array` or a tree of :class:`array`.
             in_axes (int, optional): An integer or a valid prefix tree of the
@@ -774,16 +777,16 @@ void init_transforms(py::module_& m) {
               Defaults to ``0``.
 
         Returns:
-            function: The vectorized function.
+            callable: The vectorized function.
       )pbdoc");
   m.def(
       "export_to_dot",
-      [](py::object file, const py::args& args) {
+      [](nb::object file, const nb::args& args) {
         std::vector<array> arrays = tree_flatten(args);
-        if (py::isinstance<py::str>(file)) {
-          std::ofstream out(py::cast<std::string>(file));
+        if (nb::isinstance<nb::str>(file)) {
+          std::ofstream out(nb::cast<std::string>(file));
           export_to_dot(out, arrays);
-        } else if (py::hasattr(file, "write")) {
+        } else if (nb::hasattr(file, "write")) {
           std::ostringstream out;
           export_to_dot(out, arrays);
           auto write = file.attr("write");
@@ -793,57 +796,50 @@ void init_transforms(py::module_& m) {
               "export_to_dot accepts file-like objects or strings to be used as filenames");
         }
       },
-      "file"_a);
+      "file"_a,
+      "args"_a);
   m.def(
       "compile",
-      [](const py::function& fun,
-         const py::object& inputs,
-         const py::object& outputs,
+      [](const nb::callable& fun,
+         const nb::object& inputs,
+         const nb::object& outputs,
          bool shapeless) {
-        py::options options;
-        options.disable_function_signatures();
-
-        std::ostringstream doc;
-        auto name = fun.attr("__name__").cast<std::string>();
-        doc << name;
+        //  Try to get the name
+        auto n = fun.attr("__name__");
+        auto name = n.is_none() ? "compiled" : nb::cast<std::string>(n);
 
         // Try to get the signature
-        auto inspect = py::module::import("inspect");
-        if (!inspect.attr("isbuiltin")(fun).cast<bool>()) {
-          doc << inspect.attr("signature")(fun)
-                     .attr("__str__")()
-                     .cast<std::string>();
+        std::ostringstream sig;
+        sig << "def " << name;
+        auto inspect = nb::module_::import_("inspect");
+        if (nb::cast<bool>(inspect.attr("isroutine")(fun))) {
+          sig << nb::cast<std::string>(
+              inspect.attr("signature")(fun).attr("__str__")());
+        } else {
+          sig << "(*args, **kwargs)";
         }
 
         // Try to get the doc string
-        if (auto d = fun.attr("__doc__"); py::isinstance<py::str>(d)) {
-          doc << "\n\n";
-          auto dstr = d.cast<std::string>();
-          // Add spaces to match first line indentation with remainder of
-          // docstring
-          int i = 0;
-          for (int i = dstr.size() - 1; i >= 0 && dstr[i] == ' '; i--) {
-            doc << ' ';
-          }
-          doc << dstr;
-        }
-        auto doc_str = doc.str();
-        return py::cpp_function(
+        auto d = inspect.attr("getdoc")(fun);
+        std::string doc =
+            d.is_none() ? "MLX compiled function." : nb::cast<std::string>(d);
+
+        auto sig_str = sig.str();
+        return nb::cpp_function(
             PyCompiledFun{fun, inputs, outputs, shapeless},
-            py::name(name.c_str()),
-            py::doc(doc_str.c_str()));
+            nb::name(name.c_str()),
+            nb::sig(sig_str.c_str()),
+            doc.c_str());
       },
       "fun"_a,
-      "inputs"_a = std::nullopt,
-      "outputs"_a = std::nullopt,
+      "inputs"_a = nb::none(),
+      "outputs"_a = nb::none(),
       "shapeless"_a = false,
       R"pbdoc(
-        compile(fun: function) -> function
-
         Returns a compiled function which produces the same output as ``fun``.
 
         Args:
-            fun (function): A function which takes a variable number of
+            fun (callable): A function which takes a variable number of
               :class:`array` or trees of :class:`array` and returns
               a variable number of :class:`array` or trees of :class:`array`.
             inputs (list or dict, optional): These inputs will be captured during
@@ -864,15 +860,13 @@ void init_transforms(py::module_& m) {
               ``shapeless`` set to ``True``. Default: ``False``
 
         Returns:
-            function: A compiled function which has the same input arguments
+            callable: A compiled function which has the same input arguments
             as ``fun`` and returns the the same output(s).
       )pbdoc");
   m.def(
       "disable_compile",
       &disable_compile,
       R"pbdoc(
-        disable_compile() -> None
-
         Globally disable compilation. Setting the environment variable
         ``MLX_DISABLE_COMPILE`` can also be used to disable compilation.
       )pbdoc");
@@ -880,17 +874,15 @@ void init_transforms(py::module_& m) {
       "enable_compile",
       &enable_compile,
       R"pbdoc(
-        enable_compile() -> None
-
         Globally enable compilation. This will override the environment
         variable ``MLX_DISABLE_COMPILE`` if set.
       )pbdoc");
   m.def(
       "checkpoint",
-      [](py::function fun) { return py::cpp_function(PyCheckpointedFun{fun}); },
+      [](nb::callable fun) { return nb::cpp_function(PyCheckpointedFun{fun}); },
       "fun"_a);
 
   // Register static Python object cleanup before the interpreter exits
-  auto atexit = py::module_::import("atexit");
-  atexit.attr("register")(py::cpp_function([]() { tree_cache().clear(); }));
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function([]() { tree_cache().clear(); }));
 }

python/src/trees.cpp

@@ -2,16 +2,16 @@
 
 #include "python/src/trees.h"
 
-void tree_visit(py::object tree, std::function<void(py::handle)> visitor) {
-  std::function<void(py::handle)> recurse;
-  recurse = [&](py::handle subtree) {
-    if (py::isinstance<py::list>(subtree) ||
-        py::isinstance<py::tuple>(subtree)) {
+void tree_visit(nb::object tree, std::function<void(nb::handle)> visitor) {
+  std::function<void(nb::handle)> recurse;
+  recurse = [&](nb::handle subtree) {
+    if (nb::isinstance<nb::list>(subtree) ||
+        nb::isinstance<nb::tuple>(subtree)) {
       for (auto item : subtree) {
         recurse(item);
       }
-    } else if (py::isinstance<py::dict>(subtree)) {
-      for (auto item : py::cast<py::dict>(subtree)) {
+    } else if (nb::isinstance<nb::dict>(subtree)) {
+      for (auto item : nb::cast<nb::dict>(subtree)) {
         recurse(item.second);
       }
     } else {
@@ -23,63 +23,63 @@ void tree_visit(py::object tree, std::function<void(py::handle)> visitor) {
 }
 
 template <typename T, typename U, typename V>
-void validate_subtrees(const std::vector<py::object>& subtrees) {
-  int len = py::cast<T>(subtrees[0]).size();
+void validate_subtrees(const std::vector<nb::object>& subtrees) {
+  int len = nb::cast<T>(subtrees[0]).size();
   for (auto& subtree : subtrees) {
-    if ((py::isinstance<T>(subtree) && py::cast<T>(subtree).size() != len) ||
-        py::isinstance<U>(subtree) || py::isinstance<V>(subtree)) {
+    if ((nb::isinstance<T>(subtree) && nb::cast<T>(subtree).size() != len) ||
+        nb::isinstance<U>(subtree) || nb::isinstance<V>(subtree)) {
       throw std::invalid_argument(
           "[tree_map] Additional input tree is not a valid prefix of the first tree.");
     }
   }
 }
 
-py::object tree_map(
-    const std::vector<py::object>& trees,
-    std::function<py::object(const std::vector<py::object>&)> transform) {
-  std::function<py::object(const std::vector<py::object>&)> recurse;
+nb::object tree_map(
+    const std::vector<nb::object>& trees,
+    std::function<nb::object(const std::vector<nb::object>&)> transform) {
+  std::function<nb::object(const std::vector<nb::object>&)> recurse;
 
-  recurse = [&](const std::vector<py::object>& subtrees) {
-    if (py::isinstance<py::list>(subtrees[0])) {
-      py::list l;
-      std::vector<py::object> items(subtrees.size());
-      validate_subtrees<py::list, py::tuple, py::dict>(subtrees);
-      for (int i = 0; i < py::cast<py::list>(subtrees[0]).size(); ++i) {
+  recurse = [&](const std::vector<nb::object>& subtrees) {
+    if (nb::isinstance<nb::list>(subtrees[0])) {
+      nb::list l;
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::list, nb::tuple, nb::dict>(subtrees);
+      for (int i = 0; i < nb::cast<nb::list>(subtrees[0]).size(); ++i) {
         for (int j = 0; j < subtrees.size(); ++j) {
-          if (py::isinstance<py::list>(subtrees[j])) {
-            items[j] = py::cast<py::list>(subtrees[j])[i];
+          if (nb::isinstance<nb::list>(subtrees[j])) {
+            items[j] = nb::cast<nb::list>(subtrees[j])[i];
           } else {
             items[j] = subtrees[j];
           }
         }
         l.append(recurse(items));
       }
-      return py::cast<py::object>(l);
-    } else if (py::isinstance<py::tuple>(subtrees[0])) {
+      return nb::cast<nb::object>(l);
+    } else if (nb::isinstance<nb::tuple>(subtrees[0])) {
       //  Check the rest of the subtrees
-      std::vector<py::object> items(subtrees.size());
-      int len = py::cast<py::tuple>(subtrees[0]).size();
-      py::tuple l(len);
-      validate_subtrees<py::tuple, py::list, py::dict>(subtrees);
+      std::vector<nb::object> items(subtrees.size());
+      int len = nb::cast<nb::tuple>(subtrees[0]).size();
+      nb::list l;
+      validate_subtrees<nb::tuple, nb::list, nb::dict>(subtrees);
       for (int i = 0; i < len; ++i) {
         for (int j = 0; j < subtrees.size(); ++j) {
-          if (py::isinstance<py::tuple>(subtrees[j])) {
-            items[j] = py::cast<py::tuple>(subtrees[j])[i];
+          if (nb::isinstance<nb::tuple>(subtrees[j])) {
+            items[j] = nb::cast<nb::tuple>(subtrees[j])[i];
           } else {
             items[j] = subtrees[j];
           }
         }
-        l[i] = recurse(items);
+        l.append(recurse(items));
       }
-      return py::cast<py::object>(l);
-    } else if (py::isinstance<py::dict>(subtrees[0])) {
-      std::vector<py::object> items(subtrees.size());
-      validate_subtrees<py::dict, py::list, py::tuple>(subtrees);
-      py::dict d;
-      for (auto item : py::cast<py::dict>(subtrees[0])) {
+      return nb::cast<nb::object>(nb::tuple(l));
+    } else if (nb::isinstance<nb::dict>(subtrees[0])) {
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::dict, nb::list, nb::tuple>(subtrees);
+      nb::dict d;
+      for (auto item : nb::cast<nb::dict>(subtrees[0])) {
         for (int j = 0; j < subtrees.size(); ++j) {
-          if (py::isinstance<py::dict>(subtrees[j])) {
-            auto subdict = py::cast<py::dict>(subtrees[j]);
+          if (nb::isinstance<nb::dict>(subtrees[j])) {
+            auto subdict = nb::cast<nb::dict>(subtrees[j]);
             if (!subdict.contains(item.first)) {
               throw std::invalid_argument(
                   "[tree_map] Tree is not a valid prefix tree of the first tree.");
@@ -91,7 +91,7 @@ py::object tree_map(
         }
         d[item.first] = recurse(items);
       }
-      return py::cast<py::object>(d);
+      return nb::cast<nb::object>(d);
     } else {
       return transform(subtrees);
     }
@@ -99,40 +99,40 @@ py::object tree_map(
   return recurse(trees);
 }
 
-py::object tree_map(
-    py::object tree,
-    std::function<py::object(py::handle)> transform) {
-  return tree_map({tree}, [&](std::vector<py::object> inputs) {
+nb::object tree_map(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> transform) {
+  return tree_map({tree}, [&](std::vector<nb::object> inputs) {
     return transform(inputs[0]);
   });
 }
 
 void tree_visit_update(
-    py::object tree,
-    std::function<py::object(py::handle)> visitor) {
-  std::function<py::object(py::handle)> recurse;
-  recurse = [&](py::handle subtree) {
-    if (py::isinstance<py::list>(subtree)) {
-      auto l = py::cast<py::list>(subtree);
+    nb::object tree,
+    std::function<nb::object(nb::handle)> visitor) {
+  std::function<nb::object(nb::handle)> recurse;
+  recurse = [&](nb::handle subtree) {
+    if (nb::isinstance<nb::list>(subtree)) {
+      auto l = nb::cast<nb::list>(subtree);
       for (int i = 0; i < l.size(); ++i) {
         l[i] = recurse(l[i]);
       }
-      return py::cast<py::object>(l);
-    } else if (py::isinstance<py::tuple>(subtree)) {
+      return nb::cast<nb::object>(l);
+    } else if (nb::isinstance<nb::tuple>(subtree)) {
       for (auto item : subtree) {
         recurse(item);
       }
-      return py::cast<py::object>(subtree);
-    } else if (py::isinstance<py::dict>(subtree)) {
-      auto d = py::cast<py::dict>(subtree);
+      return nb::cast<nb::object>(subtree);
+    } else if (nb::isinstance<nb::dict>(subtree)) {
+      auto d = nb::cast<nb::dict>(subtree);
       for (auto item : d) {
         d[item.first] = recurse(item.second);
       }
-      return py::cast<py::object>(d);
-    } else if (py::isinstance<array>(subtree)) {
+      return nb::cast<nb::object>(d);
+    } else if (nb::isinstance<array>(subtree)) {
       return visitor(subtree);
     } else {
-      return py::cast<py::object>(subtree);
+      return nb::cast<nb::object>(subtree);
     }
   };
   recurse(tree);
@@ -141,36 +141,36 @@ void tree_visit_update(
 // Fill a pytree (recursive dict or list of dict or list)
 // in place with the given arrays
 // Non dict or list nodes are ignored
-void tree_fill(py::object& tree, const std::vector<array>& values) {
+void tree_fill(nb::object& tree, const std::vector<array>& values) {
   size_t index = 0;
   tree_visit_update(
-      tree, [&](py::handle node) { return py::cast(values[index++]); });
+      tree, [&](nb::handle node) { return nb::cast(values[index++]); });
 }
 
 // Replace all the arrays from the src values with the dst values in the tree
 void tree_replace(
-    py::object& tree,
+    nb::object& tree,
     const std::vector<array>& src,
     const std::vector<array>& dst) {
   std::unordered_map<uintptr_t, array> src_to_dst;
   for (int i = 0; i < src.size(); ++i) {
     src_to_dst.insert({src[i].id(), dst[i]});
   }
-  tree_visit_update(tree, [&](py::handle node) {
-    auto arr = py::cast<array>(node);
+  tree_visit_update(tree, [&](nb::handle node) {
+    auto arr = nb::cast<array>(node);
     if (auto it = src_to_dst.find(arr.id()); it != src_to_dst.end()) {
-      return py::cast(it->second);
+      return nb::cast(it->second);
     }
-    return py::cast(arr);
+    return nb::cast(arr);
   });
 }
 
-std::vector<array> tree_flatten(py::object tree, bool strict /* = true */) {
+std::vector<array> tree_flatten(nb::object tree, bool strict /* = true */) {
   std::vector<array> flat_tree;
 
-  tree_visit(tree, [&](py::handle obj) {
-    if (py::isinstance<array>(obj)) {
-      flat_tree.push_back(py::cast<array>(obj));
+  tree_visit(tree, [&](nb::handle obj) {
+    if (nb::isinstance<array>(obj)) {
+      flat_tree.push_back(nb::cast<array>(obj));
     } else if (strict) {
       throw std::invalid_argument(
           "[tree_flatten] The argument should contain only arrays");
@@ -180,24 +180,24 @@ std::vector<array> tree_flatten(py::object tree, bool strict /* = true */) {
   return flat_tree;
 }
 
-py::object tree_unflatten(
-    py::object tree,
+nb::object tree_unflatten(
+    nb::object tree,
     const std::vector<array>& values,
     int index /* = 0 */) {
-  return tree_map(tree, [&](py::handle obj) {
-    if (py::isinstance<array>(obj)) {
-      return py::cast(values[index++]);
+  return tree_map(tree, [&](nb::handle obj) {
+    if (nb::isinstance<array>(obj)) {
+      return nb::cast(values[index++]);
     } else {
-      return py::cast<py::object>(obj);
+      return nb::cast<nb::object>(obj);
     }
   });
 }
 
-py::object structure_sentinel() {
-  static py::object sentinel;
+nb::object structure_sentinel() {
+  static nb::object sentinel;
 
   if (sentinel.ptr() == nullptr) {
-    sentinel = py::capsule(&sentinel);
+    sentinel = nb::capsule(&sentinel);
     // probably not needed but this should make certain that we won't ever
     // delete the sentinel
     sentinel.inc_ref();
@@ -206,19 +206,19 @@ py::object structure_sentinel() {
   return sentinel;
 }
 
-std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
-    py::object tree,
+std::pair<std::vector<array>, nb::object> tree_flatten_with_structure(
+    nb::object tree,
     bool strict /* = true */) {
   auto sentinel = structure_sentinel();
   std::vector<array> flat_tree;
   auto structure = tree_map(
       tree,
-      [&flat_tree, sentinel = std::move(sentinel), strict](py::handle obj) {
-        if (py::isinstance<array>(obj)) {
-          flat_tree.push_back(py::cast<array>(obj));
+      [&flat_tree, sentinel = std::move(sentinel), strict](nb::handle obj) {
+        if (nb::isinstance<array>(obj)) {
+          flat_tree.push_back(nb::cast<array>(obj));
           return sentinel;
         } else if (!strict) {
-          return py::cast<py::object>(obj);
+          return nb::cast<nb::object>(obj);
         } else {
           throw std::invalid_argument(
               "[tree_flatten] The argument should contain only arrays");
@@ -228,16 +228,16 @@ std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
   return {flat_tree, structure};
 }
 
-py::object tree_unflatten_from_structure(
-    py::object structure,
+nb::object tree_unflatten_from_structure(
+    nb::object structure,
     const std::vector<array>& values,
     int index /* = 0 */) {
   auto sentinel = structure_sentinel();
-  return tree_map(structure, [&](py::handle obj) {
+  return tree_map(structure, [&](nb::handle obj) {
     if (obj.is(sentinel)) {
-      return py::cast(values[index++]);
+      return nb::cast(values[index++]);
     } else {
-      return py::cast<py::object>(obj);
+      return nb::cast<nb::object>(obj);
     }
   });
 }

python/src/trees.h

@@ -1,38 +1,37 @@
 // Copyright © 2023-2024 Apple Inc.
 #pragma once
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
 
 #include "mlx/array.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 using namespace mlx::core;
 
-void tree_visit(py::object tree, std::function<void(py::handle)> visitor);
+void tree_visit(nb::object tree, std::function<void(nb::handle)> visitor);
 
-py::object tree_map(
-    const std::vector<py::object>& trees,
-    std::function<py::object(const std::vector<py::object>&)> transform);
+nb::object tree_map(
+    const std::vector<nb::object>& trees,
+    std::function<nb::object(const std::vector<nb::object>&)> transform);
 
-py::object tree_map(
-    py::object tree,
-    std::function<py::object(py::handle)> transform);
+nb::object tree_map(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> transform);
 
 void tree_visit_update(
-    py::object tree,
-    std::function<py::object(py::handle)> visitor);
+    nb::object tree,
+    std::function<nb::object(nb::handle)> visitor);
 
 /**
  * Fill a pytree (recursive dict or list of dict or list) in place with the
  * given arrays. */
-void tree_fill(py::object& tree, const std::vector<array>& values);
+void tree_fill(nb::object& tree, const std::vector<array>& values);
 
 /**
  * Replace all the arrays from the src values with the dst values in the
  * tree.
  */
 void tree_replace(
-    py::object& tree,
+    nb::object& tree,
     const std::vector<array>& src,
     const std::vector<array>& dst);
 
@@ -40,21 +39,21 @@ void tree_replace(
  * Flatten a tree into a vector of arrays. If strict is true, then the
  * function will throw if the tree contains a leaf which is not an array.
  */
-std::vector<array> tree_flatten(py::object tree, bool strict = true);
+std::vector<array> tree_flatten(nb::object tree, bool strict = true);
 
 /**
  * Unflatten a tree from a vector of arrays.
  */
-py::object tree_unflatten(
-    py::object tree,
+nb::object tree_unflatten(
+    nb::object tree,
     const std::vector<array>& values,
     int index = 0);
 
-std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
-    py::object tree,
+std::pair<std::vector<array>, nb::object> tree_flatten_with_structure(
+    nb::object tree,
     bool strict = true);
 
-py::object tree_unflatten_from_structure(
-    py::object structure,
+nb::object tree_unflatten_from_structure(
+    nb::object structure,
     const std::vector<array>& values,
     int index = 0);

python/src/utils.cpp

@@ -1,81 +0,0 @@
-
-#include "mlx/utils.h"
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
-#include <optional>
-
-namespace py = pybind11;
-using namespace py::literals;
-using namespace mlx::core;
-
-// Slightly different from the original, with python context on init we are not
-// in the context yet. Only create the inner context on enter then delete on
-// exit.
-class PyStreamContext {
- public:
-  PyStreamContext(StreamOrDevice s) : _inner(nullptr) {
-    if (std::holds_alternative<std::monostate>(s)) {
-      throw std::runtime_error(
-          "[StreamContext] Invalid argument, please specify a stream or device.");
-    }
-    _s = s;
-  }
-
-  void enter() {
-    _inner = new StreamContext(_s);
-  }
-
-  void exit() {
-    if (_inner != nullptr) {
-      delete _inner;
-      _inner = nullptr;
-    }
-  }
-
- private:
-  StreamOrDevice _s;
-  StreamContext* _inner;
-};
-
-void init_utils(py::module_& m) {
-  py::class_<PyStreamContext>(m, "StreamContext", R"pbdoc(
-        A context manager for setting the current device and stream.
-
-        See :func:`stream` for usage.
-
-        Args:
-            s: The stream or device to set as the default.
-  )pbdoc")
-      .def(py::init<StreamOrDevice>(), "s"_a)
-      .def("__enter__", [](PyStreamContext& scm) { scm.enter(); })
-      .def(
-          "__exit__",
-          [](PyStreamContext& scm,
-             const std::optional<py::type>& exc_type,
-             const std::optional<py::object>& exc_value,
-             const std::optional<py::object>& traceback) { scm.exit(); });
-  m.def(
-      "stream",
-      [](StreamOrDevice s) { return PyStreamContext(s); },
-      "s"_a,
-      R"pbdoc(
-        Create a context manager to set the default device and stream.
-
-        Args:
-            s: The :obj:`Stream` or :obj:`Device` to set as the default.
-
-        Returns:
-            A context manager that sets the default device and stream.
-
-        Example:
-
-        .. code-block::python
-
-          import mlx.core as mx
-
-          # Create a context manager for the default device and stream.
-          with mx.stream(mx.cpu):
-              # Operations here will use mx.cpu by default.
-              pass
-      )pbdoc");
-}

python/src/utils.h

@@ -1,23 +1,22 @@
-// Copyright © 2023 Apple Inc.
-
+// Copyright © 2023-2024 Apple Inc.
 #pragma once
 #include <numeric>
+#include <optional>
 #include <variant>
 
-#include <pybind11/complex.h>
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/complex.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/array.h"
 
-namespace py = pybind11;
+namespace nb = nanobind;
 
 using namespace mlx::core;
 
 using IntOrVec = std::variant<std::monostate, int, std::vector<int>>;
 using ScalarOrArray = std::
-    variant<py::bool_, py::int_, py::float_, std::complex<float>, py::object>;
-static constexpr std::monostate none{};
+    variant<nb::bool_, nb::int_, nb::float_, std::complex<float>, nb::object>;
 
 inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
   std::vector<int> axes;
@@ -32,31 +31,36 @@ inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
   return axes;
 }
 
-inline array to_array_with_accessor(py::object obj) {
-  if (py::hasattr(obj, "__mlx_array__")) {
-    return obj.attr("__mlx_array__")().cast<array>();
+inline array to_array_with_accessor(nb::object obj) {
+  if (nb::hasattr(obj, "__mlx_array__")) {
+    return nb::cast<array>(obj.attr("__mlx_array__")());
+  } else if (nb::isinstance<array>(obj)) {
+    return nb::cast<array>(obj);
   } else {
-    return obj.cast<array>();
+    std::ostringstream msg;
+    msg << "Invalid type  " << nb::type_name(obj.type()).c_str()
+        << " received in array initialization.";
+    throw std::invalid_argument(msg.str());
   }
 }
 
 inline array to_array(
     const ScalarOrArray& v,
     std::optional<Dtype> dtype = std::nullopt) {
-  if (auto pv = std::get_if<py::bool_>(&v); pv) {
-    return array(py::cast<bool>(*pv), dtype.value_or(bool_));
-  } else if (auto pv = std::get_if<py::int_>(&v); pv) {
+  if (auto pv = std::get_if<nb::bool_>(&v); pv) {
+    return array(nb::cast<bool>(*pv), dtype.value_or(bool_));
+  } else if (auto pv = std::get_if<nb::int_>(&v); pv) {
     auto out_t = dtype.value_or(int32);
     // bool_ is an exception and is always promoted
-    return array(py::cast<int>(*pv), (out_t == bool_) ? int32 : out_t);
-  } else if (auto pv = std::get_if<py::float_>(&v); pv) {
+    return array(nb::cast<int>(*pv), (out_t == bool_) ? int32 : out_t);
+  } else if (auto pv = std::get_if<nb::float_>(&v); pv) {
     auto out_t = dtype.value_or(float32);
     return array(
-        py::cast<float>(*pv), is_floating_point(out_t) ? out_t : float32);
+        nb::cast<float>(*pv), is_floating_point(out_t) ? out_t : float32);
   } else if (auto pv = std::get_if<std::complex<float>>(&v); pv) {
     return array(static_cast<complex64_t>(*pv), complex64);
   } else {
-    return to_array_with_accessor(std::get<py::object>(v));
+    return to_array_with_accessor(std::get<nb::object>(v));
   }
 }
 
@@ -68,14 +72,14 @@ inline std::pair<array, array> to_arrays(
   // - If a is an array but b is not, treat b as a weak python type
   // - If b is an array but a is not, treat a as a weak python type
   // - If neither is an array convert to arrays but leave their types alone
-  if (auto pa = std::get_if<py::object>(&a); pa) {
+  if (auto pa = std::get_if<nb::object>(&a); pa) {
     auto arr_a = to_array_with_accessor(*pa);
-    if (auto pb = std::get_if<py::object>(&b); pb) {
+    if (auto pb = std::get_if<nb::object>(&b); pb) {
       auto arr_b = to_array_with_accessor(*pb);
       return {arr_a, arr_b};
     }
     return {arr_a, to_array(b, arr_a.dtype())};
-  } else if (auto pb = std::get_if<py::object>(&b); pb) {
+  } else if (auto pb = std::get_if<nb::object>(&b); pb) {
     auto arr_b = to_array_with_accessor(*pb);
     return {to_array(a, arr_b.dtype()), arr_b};
   } else {

python/tests/test_array.py

@@ -308,9 +308,9 @@ class TestArray(mlx_tests.MLXTestCase):
         self.assertEqual(y.dtype, mx.bool_)
         self.assertEqual(y.item(), True)
 
-        # y = mx.array(x, mx.complex64)
-        # self.assertEqual(y.dtype, mx.complex64)
-        # self.assertEqual(y.item(), 3.0+0j)
+        y = mx.array(x, mx.complex64)
+        self.assertEqual(y.dtype, mx.complex64)
+        self.assertEqual(y.item(), 3.0 + 0j)
 
     def test_array_repr(self):
         x = mx.array(True)
@@ -682,7 +682,7 @@ class TestArray(mlx_tests.MLXTestCase):
 
         # check if it throws an error when dtype is not supported (bfloat16)
         x = mx.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=mx.bfloat16)
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(TypeError):
             pickle.dumps(x)
 
     def test_array_copy(self):
@@ -711,6 +711,11 @@ class TestArray(mlx_tests.MLXTestCase):
                 self.assertEqualArray(y, x - 1)
 
     def test_indexing(self):
+        # Only ellipsis is a no-op
+        a_mlx = mx.array([1])[...]
+        self.assertEqual(a_mlx.shape, (1,))
+        self.assertEqual(a_mlx.item(), 1)
+
         # Basic content check, slice indexing
         a_npy = np.arange(64, dtype=np.float32)
         a_mlx = mx.array(a_npy)
@@ -1360,7 +1365,7 @@ class TestArray(mlx_tests.MLXTestCase):
         for mlx_dtype, tf_dtype, np_dtype in dtypes_list:
             a_np = np.random.uniform(low=0, high=100, size=(3, 4)).astype(np_dtype)
             a_tf = tf.constant(a_np, dtype=tf_dtype)
-            a_mx = mx.array(a_tf)
+            a_mx = mx.array(np.array(a_tf))
             for f in [
                 lambda x: x,
                 lambda x: tf.transpose(x) if isinstance(x, tf.Tensor) else x.T,

python/tests/test_fast.py

@@ -115,6 +115,156 @@ class TestFast(mlx_tests.MLXTestCase):
                 )
                 self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
 
+    def test_rms_norm(self):
+        def rms_norm(x, weight, eps):
+            x = x.astype(mx.float32)
+            x = x * mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
+            return weight * x.astype(weight.dtype)
+
+        # Per dtype absolute tolerance
+        tolerances = {mx.float32: 1e-6, mx.float16: 1e-3, mx.bfloat16: 1e-2}
+
+        dtypes = [mx.float32, mx.float16, mx.bfloat16]
+        epss = [1e-3, 1e-5]
+        dimss = [31, 32, 33]
+        defaults = (mx.float32, 1e-5, 32)
+
+        for dtype in dtypes:
+            _, eps, dims = defaults
+            x = mx.random.uniform(
+                shape=(
+                    2,
+                    dims,
+                )
+            ).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = rms_norm(x, weight, eps)
+            rx_fast = mx.fast.rms_norm(x, weight, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        for eps in epss:
+            dtype, _, dims = defaults
+            x = mx.random.uniform(shape=(2, dims)).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = rms_norm(x, weight, eps)
+            rx_fast = mx.fast.rms_norm(x, weight, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        for dims in dimss:
+            dtype, eps, _ = defaults
+            x = mx.random.uniform(shape=(2, dims)).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = rms_norm(x, weight, eps)
+            rx_fast = mx.fast.rms_norm(x, weight, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        # Test > 4096
+        dims, dtype, eps = 4099, mx.float32, 1e-5
+        x = mx.random.uniform(shape=(dims,)).astype(dtype)
+        weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+        rx = rms_norm(x, weight, eps)
+        rx_fast = mx.fast.rms_norm(x, weight, eps)
+        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-6)
+
+    def test_layer_norm(self):
+        def layer_norm(x, weight, bias, eps):
+            ot = x.dtype
+            x = x.astype(mx.float32)
+            mean = x.mean(axis=-1, keepdims=True)
+            var = x.var(axis=-1, keepdims=True)
+            x = (x - mean) * mx.rsqrt(var + eps)
+            x = x.astype(ot)
+            if weight is not None:
+                x = x * weight
+            if bias is not None:
+                x = x + bias
+            return x
+
+        # Per dtype absolute tolerance
+        tolerances = {mx.float32: 2e-6, mx.float16: 2e-3, mx.bfloat16: 2e-2}
+
+        dtypes = [mx.float32, mx.float16, mx.bfloat16]
+        epss = [1e-3, 1e-5]
+        dimss = [31, 32, 33]
+        defaults = (mx.float32, 1e-5, 32)
+
+        for dtype in dtypes:
+            _, eps, dims = defaults
+            x = mx.random.uniform(
+                shape=(
+                    2,
+                    dims,
+                )
+            ).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            bias = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = layer_norm(x, weight, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, weight, None, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, None, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, None, eps)
+            rx_fast = mx.fast.layer_norm(x, None, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        for eps in epss:
+            dtype, _, dims = defaults
+            x = mx.random.uniform(shape=(2, dims)).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            bias = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = layer_norm(x, weight, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, weight, None, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, None, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, None, eps)
+            rx_fast = mx.fast.layer_norm(x, None, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        for dims in dimss:
+            dtype, eps, _ = defaults
+            x = mx.random.uniform(shape=(2, dims)).astype(dtype)
+            weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+            bias = mx.random.uniform(shape=(dims,)).astype(dtype)
+            rx = layer_norm(x, weight, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, weight, None, eps)
+            rx_fast = mx.fast.layer_norm(x, weight, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, bias, eps)
+            rx_fast = mx.fast.layer_norm(x, None, bias, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+            rx = layer_norm(x, None, None, eps)
+            rx_fast = mx.fast.layer_norm(x, None, None, eps)
+            self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
+        # Test > 4096
+        dims, dtype, eps = 4099, mx.float32, 1e-5
+        x = mx.random.uniform(shape=(dims,)).astype(dtype)
+        weight = mx.random.uniform(shape=(dims,)).astype(dtype)
+        bias = mx.random.uniform(shape=(dims,)).astype(dtype)
+        rx = layer_norm(x, weight, bias, eps)
+        rx_fast = mx.fast.layer_norm(x, weight, bias, eps)
+        self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+        rx = layer_norm(x, weight, None, eps)
+        rx_fast = mx.fast.layer_norm(x, weight, None, eps)
+        self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+        rx = layer_norm(x, None, bias, eps)
+        rx_fast = mx.fast.layer_norm(x, None, bias, eps)
+        self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+        rx = layer_norm(x, None, None, eps)
+        rx_fast = mx.fast.layer_norm(x, None, None, eps)
+        self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
+
     def test_fast_transforms(self):
         x = mx.random.uniform(shape=(2, 2, 8))
 

python/tests/test_linalg.py

@@ -136,6 +136,20 @@ class TestLinalg(mlx_tests.MLXTestCase):
                 mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
             )
 
+    def test_inverse(self):
+        A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32)
+        A_inv = mx.linalg.inv(A, stream=mx.cpu)
+        self.assertTrue(mx.allclose(A @ A_inv, mx.eye(A.shape[0]), rtol=0, atol=1e-6))
+
+        # Multiple matrices
+        B = A - 100
+        AB = mx.stack([A, B])
+        invs = mx.linalg.inv(AB, stream=mx.cpu)
+        for M, M_inv in zip(AB, invs):
+            self.assertTrue(
+                mx.allclose(M @ M_inv, mx.eye(M.shape[0]), rtol=0, atol=1e-5)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_ops.py

@@ -590,6 +590,18 @@ class TestOps(mlx_tests.MLXTestCase):
             self.assertListEqual(list(sum_npy.shape), list(sum_mlx.shape))
             self.assertTrue(np.all(sum_npy == sum_mlx))
 
+        x_npy = (
+            np.arange(3 * 2 * 3 * 3 * 3 * 3)
+            .reshape(3, 2, 3, 3, 3, 3)
+            .astype(np.float32)
+        )
+        x_mlx = mx.array(x_npy)
+
+        y_mlx = x_mlx.sum(axis=(0, 1, 3, 4, 5))
+        y_npy = x_npy.sum(axis=(0, 1, 3, 4, 5))
+
+        self.assertTrue(np.array_equal(y_mlx, y_npy))
+
     def test_prod(self):
         x = mx.array(
             [

python/tests/test_optimizers.py

@@ -328,6 +328,11 @@ class TestSchedulers(unittest.TestCase):
         expected_lr = 0.1 * 0.5 * (1.0 + math.cos(math.pi * 4 / 10))
         self.assertAlmostEqual(lr, expected_lr, delta=1e-7)
 
+        lr_schedule = opt.cosine_decay(0.1, 10, 0.05)
+        lr = lr_schedule(20)
+        expected_lr = 0.05
+        self.assertEqual(lr, expected_lr)
+
     def test_schedule_joiner(self):
         boundaries = [2, 3, 4]
         schedules = [lambda _: 3, lambda _: 4, lambda _: 5]

python/tests/test_vmap.py

@@ -275,6 +275,103 @@ class TestVmap(mlx_tests.MLXTestCase):
         with self.assertRaises(ValueError):
             out = mx.vmap(lambda x, y: x + y, in_axes=(0, 1))(a, b)
 
+    def test_vmap_matmul(self):
+        a = mx.random.uniform(shape=(2, 3, 4))
+        b = mx.random.uniform(shape=(4, 3))
+
+        # matmul
+        out = mx.vmap(mx.matmul, in_axes=(0, None))(a, b)
+        self.assertTrue(mx.allclose(out, a @ b))
+
+        # addmm
+        c = mx.random.uniform(shape=(3,))
+        out = mx.vmap(mx.addmm, in_axes=(None, 0, None))(c, a, b)
+        self.assertTrue(mx.allclose(out, mx.addmm(c, a, b)))
+
+        b = mx.random.uniform(shape=(4, 2))
+
+        # matmul
+        out = mx.vmap(mx.matmul, in_axes=(1, None), out_axes=(1,))(a, b)
+        expected = mx.moveaxis(mx.moveaxis(a, 1, 0) @ b, 0, 1)
+        self.assertTrue(mx.allclose(out, expected))
+
+        # addmm
+        c = mx.random.uniform(shape=(2,))
+        out = mx.vmap(mx.addmm, in_axes=(None, 1, None))(c, a, b)
+        self.assertTrue(mx.allclose(out, mx.addmm(c, mx.moveaxis(a, 1, 0), b)))
+
+        a = mx.random.uniform(shape=(2, 3, 4))
+        b = mx.random.uniform(shape=(4, 2, 3))
+
+        # matmul
+        out = mx.vmap(mx.matmul, in_axes=(0, 1))(a, b)
+        expected = a @ mx.moveaxis(b, 1, 0)
+        self.assertTrue(mx.allclose(out, expected))
+
+        # addmm
+        c = mx.random.uniform(shape=(3, 3, 2))
+        out = mx.vmap(mx.addmm, in_axes=(2, 0, 1))(c, a, b)
+        expected = mx.addmm(mx.moveaxis(c, 2, 0), a, mx.moveaxis(b, 1, 0))
+        self.assertTrue(mx.allclose(out, expected))
+
+    def test_vmap_svd(self):
+        a = mx.random.uniform(shape=(3, 4, 2))
+
+        cpu_svd = lambda x: mx.linalg.svd(x, stream=mx.cpu)
+
+        # Vmap over the first axis (this is already supported natively by the primitive).
+        Us, Ss, Vts = mx.vmap(cpu_svd, in_axes=(0,))(a)
+        self.assertEqual(Us.shape, (a.shape[0], a.shape[1], a.shape[1]))
+        self.assertEqual(Ss.shape, (a.shape[0], a.shape[2]))
+        self.assertEqual(Vts.shape, (a.shape[0], a.shape[2], a.shape[2]))
+
+        for i in range(a.shape[0]):
+            M = a[i]
+            U, S, Vt = Us[i], Ss[i], Vts[i]
+            self.assertTrue(
+                mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
+            )
+
+        # Vmap over the second axis.
+        Us, Ss, Vts = mx.vmap(cpu_svd, in_axes=(1,))(a)
+        self.assertEqual(Us.shape, (a.shape[1], a.shape[0], a.shape[0]))
+        self.assertEqual(Ss.shape, (a.shape[1], a.shape[2]))
+        self.assertEqual(Vts.shape, (a.shape[1], a.shape[2], a.shape[2]))
+
+        for i in range(a.shape[1]):
+            M = a[:, i, :]
+            U, S, Vt = Us[i], Ss[i], Vts[i]
+            self.assertTrue(
+                mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
+            )
+
+    def test_vmap_inverse(self):
+        a = mx.random.uniform(shape=(3, 4, 4))
+
+        cpu_inv = lambda x: mx.linalg.inv(x, stream=mx.cpu)
+
+        # Vmap over the first axis (this is already supported natively by the primitive).
+        invs = mx.vmap(cpu_inv, in_axes=(0,))(a)
+
+        for i in range(a.shape[0]):
+            self.assertTrue(
+                mx.allclose(a[i] @ invs[i], mx.eye(a.shape[1]), rtol=0, atol=1e-5)
+            )
+
+        a = mx.random.uniform(shape=(4, 3, 4))
+
+        # Without vmapping, each input matrix is not square.
+        with self.assertRaises(ValueError):
+            mx.eval(cpu_inv(a))
+
+        # Vmap over the second axis.
+        invs = mx.vmap(cpu_inv, in_axes=(1,))(a)
+
+        for i in range(a.shape[1]):
+            self.assertTrue(
+                mx.allclose(a[:, i, :] @ invs[i], mx.eye(a.shape[0]), rtol=0, atol=1e-5)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

setup.py

@@ -134,7 +134,9 @@ class GenerateStubs(Command):
         pass
 
     def run(self) -> None:
-        subprocess.run(["pybind11-stubgen", "mlx.core", "-o", "python"])
+        subprocess.run(
+            ["python", "-m", "nanobind.stubgen", "-m", "mlx.core", "-r", "-O", "python"]
+        )
 
 
 # Read the content of README.md
@@ -152,7 +154,7 @@ if __name__ == "__main__":
 
     setup(
         name="mlx",
-        version=get_version("0.7.0"),
+        version=get_version("0.8.0"),
         author="MLX Contributors",
         author_email="mlx@group.apple.com",
         description="A framework for machine learning on Apple silicon.",
@@ -165,7 +167,7 @@ if __name__ == "__main__":
         include_package_data=True,
         extras_require={
             "testing": ["numpy", "torch"],
-            "dev": ["pre-commit", "pybind11-stubgen"],
+            "dev": ["pre-commit"],
         },
         ext_modules=[CMakeExtension("mlx.core")],
         cmdclass={"build_ext": CMakeBuild, "generate_stubs": GenerateStubs},

tests/linalg_tests.cpp

@@ -300,3 +300,25 @@ TEST_CASE("test SVD factorization") {
   CHECK_EQ(S.dtype(), float32);
   CHECK_EQ(Vt.dtype(), float32);
 }
+
+TEST_CASE("test matrix inversion") {
+  // 0D and 1D throw
+  CHECK_THROWS(linalg::inv(array(0.0), Device::cpu));
+  CHECK_THROWS(linalg::inv(array({0.0, 1.0}), Device::cpu));
+
+  // Unsupported types throw
+  CHECK_THROWS(linalg::inv(array({0, 1}, {1, 2}), Device::cpu));
+
+  // Non-square throws.
+  CHECK_THROWS(linalg::inv(array({1, 2, 3, 4, 5, 6}, {2, 3}), Device::cpu));
+
+  const auto prng_key = random::key(42);
+  const auto A = random::normal({5, 5}, prng_key);
+  const auto A_inv = linalg::inv(A, Device::cpu);
+  const auto identity = eye(A.shape(0));
+
+  CHECK(allclose(matmul(A, A_inv), identity, /* rtol = */ 0, /* atol = */ 1e-6)
+            .item<bool>());
+  CHECK(allclose(matmul(A_inv, A), identity, /* rtol = */ 0, /* atol = */ 1e-6)
+            .item<bool>());
+}

tests/ops_tests.cpp

@@ -198,6 +198,31 @@ TEST_CASE("test slice") {
   CHECK(array_equal(out, array({0, 2, 4, 6}, {2, 2})).item<bool>());
 }
 
+TEST_CASE("test slice update") {
+  array x = array({0., 0., 0., 0., 0., 0., 0., 0.}, {8}, float32);
+  array y = array(
+      {
+          1.,
+          2.,
+          3.,
+          4.,
+      },
+      {4},
+      float32);
+
+  auto out = slice_update(x, y, {2}, {6}, {1});
+  CHECK(array_equal(slice(out, {2}, {6}, {1}), y).item<bool>());
+
+  out = slice_update(x, y, {5}, {1}, {-1});
+  CHECK(array_equal(slice(out, {5}, {1}, {-1}), y).item<bool>());
+
+  x = reshape(x, {2, 4});
+  out = slice_update(x, y, {0, 0}, {2, 4}, {1, 1});
+  out = reshape(out, {8});
+  CHECK(array_equal(slice(out, {0}, {4}, {1}), y).item<bool>());
+  CHECK(array_equal(slice(out, {4}, {8}, {1}), y).item<bool>());
+}
+
 TEST_CASE("test split") {
   array x = array(1);
   CHECK_THROWS(split(x, 0));

tests/vmap_tests.cpp

@@ -413,3 +413,35 @@ TEST_CASE("test vmap gather") {
     CHECK_EQ(out.shape(), std::vector<int>{2, 3, 2, 2});
   }
 }
+
+TEST_CASE("test vmap SVD") {
+  auto fun = [](std::vector<array> inputs) {
+    return linalg::svd(inputs.at(0), Device::cpu);
+  };
+
+  auto a = astype(reshape(arange(24), {3, 4, 2}), float32);
+
+  // vmap over the second axis.
+  {
+    auto out = vmap(fun, /* in_axes = */ {1})({a});
+    const auto& U = out.at(0);
+    const auto& S = out.at(1);
+    const auto& Vt = out.at(2);
+
+    CHECK_EQ(U.shape(), std::vector<int>{a.shape(1), a.shape(0), a.shape(0)});
+    CHECK_EQ(S.shape(), std::vector<int>{a.shape(1), a.shape(2)});
+    CHECK_EQ(Vt.shape(), std::vector<int>{a.shape(1), a.shape(2), a.shape(2)});
+  }
+
+  // vmap over the third axis.
+  {
+    auto out = vmap(fun, /* in_axes = */ {2})({a});
+    const auto& U = out.at(0);
+    const auto& S = out.at(1);
+    const auto& Vt = out.at(2);
+
+    CHECK_EQ(U.shape(), std::vector<int>{a.shape(2), a.shape(0), a.shape(0)});
+    CHECK_EQ(S.shape(), std::vector<int>{a.shape(2), a.shape(0)});
+    CHECK_EQ(Vt.shape(), std::vector<int>{a.shape(2), a.shape(1), a.shape(1)});
+  }
+}