try again with checkpointed classes

.circleci/config.yml

@@ -31,7 +31,8 @@ jobs:
           name: Install dependencies
           command: |
             pip install --upgrade cmake
-            pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+            pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             sudo apt-get update
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
@@ -43,12 +44,16 @@ jobs:
       - run:
           name: Generate package stubs
           command: |
-            echo "stubs"
-            python setup.py generate_stubs 
+            python3 setup.py generate_stubs
       - run:
           name: Run Python tests
           command: |
             python3 -m unittest discover python/tests -v
+      # TODO: Reenable when extension api becomes stable
+      # - run:
+      #     name: Build example extension
+      #     command: |
+      #       cd examples/extensions && python3 -m pip install . 
       - run:
           name: Build CPP only
           command: |
@@ -58,25 +63,21 @@ jobs:
           command: ./build/tests/tests
 
   mac_build_and_test:
-    parameters:
-      xcode_version:
-        type: string
-        default: "15.2.0"
     macos:
-      xcode: << parameters.xcode_version >>
-    resource_class: macos.m1.medium.gen1
+      xcode: "15.2.0"
+    resource_class: macos.m1.large.gen1
     steps:
       - checkout
       - run:
           name: Install dependencies
           command: |
-            brew install python@3.8
-            brew install openmpi
-            python3.8 -m venv env
+            brew install python@3.9
+            python3.9 -m venv env
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+            pip install --upgrade pybind11[global]
+            pip install pybind11-stubgen
             pip install numpy
             pip install torch
             pip install tensorflow
@@ -90,21 +91,18 @@ jobs:
           name: Generate package stubs
           command: |
             source env/bin/activate
-            python setup.py generate_stubs 
+            python setup.py generate_stubs
       - run:
           name: Run Python tests
           command: |
             source env/bin/activate
             LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
-            mpirun -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
-      - run:
-          name: Build example extension
-          command: |
-            source env/bin/activate
-            cd examples/extensions
-            pip install -r requirements.txt
-            python setup.py build_ext -j8
+            LOW_MEMORY=1 DEVICE=gpu python3.9 -m xmlrunner discover -v python/tests -o test-results/gpu
+      # TODO: Reenable when extension api becomes stable
+      # - run:
+      #     name: Build example extension
+      #     command: |
+      #       cd examples/extensions && python3.11 -m pip install . 
       - store_test_results:
           path: test-results
       - run:
@@ -116,13 +114,7 @@ jobs:
           name: Run CPP tests
           command: |
             DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
-      - run:
-          name: Build small binary
-          command: |
-            source env/bin/activate
-            cd build/
-            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel -DBUILD_SHARED_LIBS=ON -DMLX_BUILD_CPU=OFF -DMLX_BUILD_SAFETENSORS=OFF -DMLX_BUILD_GGUF=OFF -DMLX_METAL_JIT=ON
-            make -j
+            DEVICE=cpu ./build/tests/tests
 
   build_release:
     parameters:
@@ -137,20 +129,20 @@ jobs:
         default: ""
     macos:
       xcode: << parameters.xcode_version >>
-    resource_class: macos.m1.medium.gen1
+    resource_class: macos.m1.large.gen1
     steps:
       - checkout
       - run:
           name: Install dependencies
           command: |
             brew install python@<< parameters.python_version >>
-            brew install openmpi
             python<< parameters.python_version >> -m venv env
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+            pip install --upgrade pybind11[global]
             pip install --upgrade setuptools
+            pip install pybind11-stubgen
             pip install numpy
             pip install twine
             pip install build
@@ -165,7 +157,7 @@ jobs:
           name: Generate package stubs
           command: |
             source env/bin/activate
-            python setup.py generate_stubs 
+            python setup.py generate_stubs
       - run:
           name: Build Python package
           command: |
@@ -213,8 +205,9 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+            pip install --upgrade pybind11[global]
             pip install --upgrade setuptools
+            pip install pybind11-stubgen
             pip install numpy
             pip install auditwheel
             pip install patchelf
@@ -222,7 +215,7 @@ jobs:
             << parameters.extra_env >> \
               CMAKE_BUILD_PARALLEL_LEVEL="" \
               pip install . -v
-            python setup.py generate_stubs 
+            python setup.py generate_stubs
             << parameters.extra_env >> \
               CMAKE_BUILD_PARALLEL_LEVEL="" \
               python -m build --wheel
@@ -242,10 +235,7 @@ workflows:
         - not: << pipeline.parameters.weekly_build >>
         - not: << pipeline.parameters.test_release >>
     jobs:
-      - mac_build_and_test:
-          matrix:
-            parameters:
-              xcode_version: ["15.0.0", "15.2.0"]
+      - mac_build_and_test
       - linux_build_and_test
 
   build_pypi_release:
@@ -264,7 +254,7 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-              xcode_version: ["15.0.0", "15.2.0"]
+              xcode_version: ["14.3.1", "15.2.0"]
               build_env: ["PYPI_RELEASE=1"]
   prb:
     when:
@@ -278,9 +268,6 @@ workflows:
           context: pr-approval
       - mac_build_and_test:
           requires: [ hold ]
-          matrix:
-            parameters:
-              xcode_version: ["15.0.0", "15.2.0"]
       - linux_build_and_test:
           requires: [ hold ]
   nightly_build:
@@ -293,7 +280,7 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-              xcode_version: ["15.0.0", "15.2.0"]
+              xcode_version: ["14.3.1", "15.2.0"]
   weekly_build:
     when:
       and:
@@ -304,7 +291,7 @@ workflows:
           matrix:
             parameters:
               python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-              xcode_version: ["15.0.0", "15.2.0"]
+              xcode_version: ["14.3.1", "15.2.0"]
               build_env: ["DEV_RELEASE=1"]
   linux_test_release:
     when:

.github/workflows/pull_request.yml

@@ -17,4 +17,4 @@ jobs:
           pip install pre-commit black isort clang-format
       - name: Run lint
         run: |
-          pre-commit run --all-files
+          pre-commit run --all-files

.pre-commit-config.yaml

@@ -1,11 +1,11 @@
 repos:
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.4
+    rev: v17.0.6
     hooks:
     -   id: clang-format
 # Using this mirror lets us use mypyc-compiled black, which is about 2x faster
 -   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 24.4.2
+    rev: 24.2.0
     hooks:
     -   id: black
 -   repo: https://github.com/pycqa/isort

ACKNOWLEDGMENTS.md

@@ -7,18 +7,14 @@ with a short description of your contribution(s) below. For example:
 
 MLX was developed with contributions from the following individuals:
 
-- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`.
+- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops.
 - Juarez Bochi: Fixed bug in cross attention.
 - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example.
-- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`.
+- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream` and safetensor support.
 - Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented pooling layers and ``Upsample``.
 - Hinrik Snær Guðmundsson: Added `atleast_1d`, `atleast_2d`, `atleast_3d` ops.
 - Luca Arnaboldi: Added `Ceil` and `Floor` ops; implemented pickling, copy and deepcopy for mlx arrays.
 - Brian Keene & Atila Orhon, with Argmax Inc.: Added `fast.scaled_dot_product_attention`
-- AmirHossein Razlighi: Added chaining support for some of the ops in `nn.Module`. Comparison works for non array objects in `mlx.core.array`. Exception handling for invalid operations in `mlx.core.array`.
-- Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
-- Paul Paczuski: Improved stability of BCE loss calculation
-
 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
 </a>

CMakeLists.txt

@@ -15,37 +15,31 @@ option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON)
 option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
 option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
 option(MLX_BUILD_METAL "Build metal backend" ON)
-option(MLX_BUILD_CPU "Build cpu backend" ON)
-option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
-option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
-option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
-option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
-option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.16.3)
+  set(MLX_VERSION 0.5.1)
 endif()
 
 # --------------------- Processor tests -------------------------
 
-message(STATUS "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}")
+message(STATUS "Building MLX for ${CMAKE_HOST_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}")
 
 set(MLX_BUILD_ARM OFF)
 
 if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
-  if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64")
-    if(NOT MLX_ENABLE_X64_MAC)
-      message(FATAL_ERROR
-        "Building for x86_64 on macOS is not supported."
-        " If you are on an Apple silicon system, check the build"
-        " documentation for possible fixes: "
-        "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source")
-    else()
-      message(WARNING "Building for x86_64 arch is not officially supported.")
-    endif()
-    set(MLX_BUILD_METAL OFF)
-  elseif(${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64")
+  if (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64" AND ${CMAKE_HOST_APPLE})
+    message(FATAL_ERROR
+      "Building for x86_64 on macOS is not supported."
+      " If you are on an Apple silicon system, check the build"
+      " documentation for possible fixes: "
+      "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source")
+  elseif (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64")
+    message(WARNING
+      "Building for x86_64 on macOS is not supported."
+      " If you are on an Apple silicon system, "
+      " make sure you are building for arm64.")
+  elseif(${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "arm64")
     set(MLX_BUILD_ARM ON)
   endif()
 
@@ -70,30 +64,24 @@ endif()
 if (MLX_BUILD_METAL AND NOT METAL_LIB)
   message(STATUS "Metal not found. Unable to build GPU")
   set(MLX_BUILD_METAL OFF)
-  set(MLX_METAL_DEBUG OFF)
 elseif (MLX_BUILD_METAL)
   message(STATUS "Building METAL sources")
-
-  if (MLX_METAL_DEBUG)
-    add_compile_definitions(MLX_METAL_DEBUG)
-  endif()
-
   # Throw an error if xcrun not found
   execute_process(COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
                   OUTPUT_VARIABLE MACOS_VERSION
                   COMMAND_ERROR_IS_FATAL ANY)
 
-  if (${MACOS_VERSION} LESS 14.0)
-    message(FATAL_ERROR "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON" )
-  endif()
   message(STATUS "Building with SDK for macOS version ${MACOS_VERSION}")
 
-  set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18-beta.zip)
-  # Get the metal version
-  execute_process(
-    COMMAND zsh "-c" "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal -E -x metal -P - | tail -1 | tr -d '\n'"
-    OUTPUT_VARIABLE MLX_METAL_VERSION
-    COMMAND_ERROR_IS_FATAL ANY)
+  if (${MACOS_VERSION} GREATER_EQUAL 14.2)
+    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14.2_iOS17.2.zip)
+  elseif (${MACOS_VERSION} GREATER_EQUAL 14.0)
+    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14_iOS17-beta.zip)
+  elseif (${MACOS_VERSION} GREATER_EQUAL 13.3)
+    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS13.3_iOS16.4.zip)
+  else()
+    message(FATAL_ERROR "MLX requires macOS >= 13.4 to be built with MLX_BUILD_METAL=ON" )
+  endif()
 
   FetchContent_Declare(
     metal_cpp
@@ -107,85 +95,46 @@ elseif (MLX_BUILD_METAL)
     $<INSTALL_INTERFACE:include/metal_cpp>
   )
   target_link_libraries(
-    mlx PUBLIC
+    mlx
     ${METAL_LIB}
     ${FOUNDATION_LIB}
     ${QUARTZ_LIB})
-
-  add_compile_definitions("MLX_METAL_VERSION=${MLX_METAL_VERSION}")
 endif()
 
-if (MLX_BUILD_CPU)
-  find_library(ACCELERATE_LIBRARY Accelerate)
-  if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY)
-    message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
-    set(MLX_BUILD_ACCELERATE ON)
-    target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
-    add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  else()
-    message(STATUS "Accelerate or arm neon not found, using default backend.")
-    set(MLX_BUILD_ACCELERATE OFF)
-    if(${CMAKE_HOST_APPLE})
-      # The blas shipped in macOS SDK is not supported, search homebrew for
-      # openblas instead.
-      set(BLA_VENDOR OpenBLAS)
-      set(LAPACK_ROOT "${LAPACK_ROOT};$ENV{LAPACK_ROOT};/usr/local/opt/openblas")
-    endif()
-    # Search and link with lapack.
-    find_package(LAPACK REQUIRED)
-    if (NOT LAPACK_FOUND)
-      message(FATAL_ERROR "Must have LAPACK installed")
-    endif()
-    find_path(LAPACK_INCLUDE_DIRS lapacke.h
-      /usr/include
-      /usr/local/include
-      /usr/local/opt/openblas/include)
-    message(STATUS "Lapack lib " ${LAPACK_LIBRARIES})
-    message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS})
-    target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS})
-    target_link_libraries(mlx PUBLIC ${LAPACK_LIBRARIES})
-    # List blas after lapack otherwise we may accidentally incldue an old version
-    # of lapack.h from the include dirs of blas.
-    find_package(BLAS REQUIRED)
-    if (NOT BLAS_FOUND)
-      message(FATAL_ERROR "Must have BLAS installed")
-    endif()
-    # TODO find a cleaner way to do this
-    find_path(BLAS_INCLUDE_DIRS cblas.h
-      /usr/include
-      /usr/local/include
-      $ENV{BLAS_HOME}/include)
-    message(STATUS "Blas lib " ${BLAS_LIBRARIES})
-    message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS})
-    target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS})
-    target_link_libraries(mlx PUBLIC ${BLAS_LIBRARIES})
-  endif()
+find_library(ACCELERATE_LIBRARY Accelerate)
+if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY)
+  message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
+  set(MLX_BUILD_ACCELERATE ON)
+  target_link_libraries(mlx ${ACCELERATE_LIBRARY})
+  add_compile_definitions(ACCELERATE_NEW_LAPACK)
 else()
+  message(STATUS "Accelerate or arm neon not found, using default backend.")
   set(MLX_BUILD_ACCELERATE OFF)
-endif()
-
-find_package(MPI)
-if (MPI_FOUND)
-  execute_process(
-    COMMAND zsh "-c" "mpirun --version"
-    OUTPUT_VARIABLE MPI_VERSION
-    ERROR_QUIET
-  )
-  if (${MPI_VERSION} MATCHES ".*Open MPI.*")
-    target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
-  elseif (MPI_VERSION STREQUAL "")
-    set(MPI_FOUND FALSE)
-    message(
-      WARNING
-      "MPI found but mpirun is not available. Building without MPI."
-    )
-  else()
-    set(MPI_FOUND FALSE)
-    message(
-      WARNING
-      "MPI which is not OpenMPI found. Building without MPI."
-    )
-  endif() 
+  #set(BLA_VENDOR Generic)
+  find_package(BLAS REQUIRED)
+  if (NOT BLAS_FOUND)
+    message(FATAL_ERROR "Must have BLAS installed")
+  endif()
+  # TODO find a cleaner way to do this
+  find_path(BLAS_INCLUDE_DIRS cblas.h
+    /usr/include
+    /usr/local/include
+    $ENV{BLAS_HOME}/include)
+  message(STATUS "Blas lib " ${BLAS_LIBRARIES})
+  message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS})
+  target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS})
+  target_link_libraries(mlx ${BLAS_LIBRARIES})
+  find_package(LAPACK REQUIRED)
+  if (NOT LAPACK_FOUND)
+      message(FATAL_ERROR "Must have LAPACK installed")
+  endif()
+  find_path(LAPACK_INCLUDE_DIRS lapacke.h
+    /usr/include
+    /usr/local/include)
+  message(STATUS "Lapack lib " ${LAPACK_LIBRARIES})
+  message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS})
+  target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS})
+  target_link_libraries(mlx ${LAPACK_LIBRARIES})
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
@@ -197,22 +146,10 @@ target_include_directories(
   $<INSTALL_INTERFACE:include>
 )
 
-FetchContent_Declare(fmt
-  GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-  GIT_TAG 10.2.1 
-  EXCLUDE_FROM_ALL
-)
-FetchContent_MakeAvailable(fmt)
-target_link_libraries(mlx PRIVATE fmt::fmt-header-only)
-
 if (MLX_BUILD_PYTHON_BINDINGS)
   message(STATUS "Building Python bindings.")
-  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)
+  find_package(Python COMPONENTS Interpreter Development)
+  find_package(pybind11 CONFIG REQUIRED)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 

README.md

@@ -88,13 +88,13 @@ for more information on building the C++ and Python APIs from source.
 
 ## Contributing 
 
-Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
+Check out the [contribution guidelines](CONTRIBUTING.md) for more information
 on contributing to MLX. See the
 [docs](https://ml-explore.github.io/mlx/build/html/install.html) for more
 information on building from source, and running tests.
 
 We are grateful for all of [our
-contributors](https://github.com/ml-explore/mlx/tree/main/ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
+contributors](ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
 to MLX and wish to be acknowledged, please add your name to the list in your
 pull request.
 

benchmarks/cpp/time_utils.h

@@ -17,13 +17,14 @@
             << std::setprecision(5) << time_fn(FUNC, ##__VA_ARGS__) << " msec" \
             << std::endl;
 
-#define TIMEM(MSG, FUNC, ...)                                      \
-  std::cout << "Timing " << "(" << MSG << ") " << #FUNC << " ... " \
-            << std::flush << std::setprecision(5)                  \
-            << time_fn(FUNC, ##__VA_ARGS__) << " msec" << std::endl;
+#define TIMEM(MSG, FUNC, ...)                                                  \
+  std::cout << "Timing "                                                       \
+            << "(" << MSG << ") " << #FUNC << " ... " << std::flush            \
+            << std::setprecision(5) << time_fn(FUNC, ##__VA_ARGS__) << " msec" \
+            << std::endl;
 
 template <typename F, typename... Args>
-double time_fn(F fn, Args&&... args) {
+double time_fn(F fn, Args... args) {
   // warmup
   for (int i = 0; i < 5; ++i) {
     eval(fn(std::forward<Args>(args)...));

benchmarks/python/comparative/bench_torch.py

@@ -185,7 +185,7 @@ def prelu(x: torch.Tensor) -> torch.Tensor:
 def mish(x: torch.Tensor) -> torch.Tensor:
     y = x
     for _ in range(100):
-        y = torch.nn.functional.mish(y)
+        return torch.nn.functional.mish(y)
     sync_if_needed(x)
 
 
@@ -283,14 +283,6 @@ def topk(axis, x):
     sync_if_needed(x)
 
 
-@torch.no_grad()
-def step_function(x):
-    y = x
-    for i in range(100):
-        y = torch.where(y < 0, 0, 1)
-    sync_if_needed(x)
-
-
 @torch.no_grad()
 def selu(x):
     y = x
@@ -454,11 +446,5 @@ if __name__ == "__main__":
     elif args.benchmark == "topk":
         print(bench(topk, axis, x))
 
-    elif args.benchmark == "step":
-        print(bench(step_function, x))
-
-    elif args.benchmark == "selu":
-        print(bench(selu, x))
-
     else:
-        raise ValueError(f"Unknown benchmark `{args.benchmark}`.")
+        raise ValueError("Unknown benchmark")

benchmarks/python/comparative/compare.py

@@ -16,9 +16,7 @@ def run_or_raise(*args, **kwargs):
         result = run(*args, capture_output=True, **kwargs)
         return float(result.stdout)
     except ValueError:
-        raise ValueError(
-            f"stdout: {result.stdout.decode()}\nstderr: {result.stderr.decode()}"
-        )
+        raise ValueError(f"stdout: {result.stdout}\nstderr: {result.stderr}")
 
 
 def compare(args):

benchmarks/python/compile_bench.py

@@ -9,6 +9,7 @@ from time_utils import time_fn
 
 
 def bench_gelu():
+
     def gelu(x):
         return x * (1 + mx.erf(x / math.sqrt(2))) / 2
 
@@ -50,6 +51,7 @@ def bench_gelu():
 
 
 def bench_layernorm():
+
     weight = mx.random.uniform(shape=(4096,)).astype(mx.float16)
     bias = mx.random.uniform(shape=(4096,)).astype(mx.float16)
     mx.eval(weight, bias)

benchmarks/python/conv1d_bench.py

@@ -1,123 +0,0 @@
-import argparse
-import math
-import os
-import subprocess
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
-device_name = device_name.decode("utf-8").strip("\n")
-
-N_warmup = 10
-N_iter_bench = 100
-N_iter_func = 5
-
-
-def bench(f, a, b):
-    for i in range(N_warmup):
-        f(a, b)
-    torch.mps.synchronize()
-
-    s = time.perf_counter_ns()
-    for i in range(N_iter_bench):
-        f(a, b)
-    e = time.perf_counter_ns()
-    return (e - s) * 1e-9
-
-
-def make_mx_conv_1D(strides=1, padding=0, groups=1):
-    def mx_conv_1D(a, b):
-        ys = []
-        for _ in range(N_iter_func):
-            y = mx.conv1d(a, b, stride=strides, padding=padding, groups=groups)
-            ys.append(y)
-        mx.eval(ys)
-        return ys
-
-    return mx_conv_1D
-
-
-def make_pt_conv_1D(strides=1, padding=0, groups=1):
-    @torch.no_grad()
-    def pt_conv_1D(a, b):
-        ys = []
-        for _ in range(N_iter_func):
-            y = torch.conv1d(a, b, stride=strides, padding=padding, groups=groups)
-            ys.append(y)
-        torch.mps.synchronize()
-        return ys
-
-    return pt_conv_1D
-
-
-def bench_shape(N, iH, C, wH, O, strides, padding, np_dtype, groups):
-    scale = 1.0 / math.sqrt(wH * C)
-    a_np = np.random.uniform(0, 0.5, (N, iH, C)).astype(np_dtype)
-    b_np = np.random.uniform(-scale, scale, (O, wH, int(C / groups))).astype(np_dtype)
-
-    a_mx = mx.array(a_np)
-    b_mx = mx.array(b_np)
-
-    a_pt = torch.from_numpy(a_np.transpose((0, 2, 1))).to("mps")
-    b_pt = torch.from_numpy(b_np.transpose((0, 2, 1))).to("mps")
-
-    torch.mps.synchronize()
-
-    f_mx = make_mx_conv_1D(strides, padding, groups)
-    f_pt = make_pt_conv_1D(strides, padding, groups)
-
-    time_torch = bench(f_pt, a_pt, b_pt)
-    time_mlx = bench(f_mx, a_mx, b_mx)
-
-    out_mx = mx.conv1d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
-    out_pt = torch.conv1d(
-        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
-    )
-    out_pt = torch.permute(out_pt, (0, 2, 1))
-    out_pt = out_pt.numpy(force=True)
-
-    atol = 2e-5 if np_dtype == np.float32 else 1e-4
-
-    if not np.allclose(out_pt, out_mx, atol=atol):
-        print(
-            f"Failed at {(N, iH, C)}, {(O, wH, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
-        )
-
-    return time_mlx, time_torch
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run conv benchmarks")
-
-    dtypes = ("float32",)
-    shapes = (
-        (4, 32, 32, 5, 32, 1, 2, 1),
-        (4, 32, 32, 5, 32, 1, 2, 2),
-        (4, 32, 32, 5, 32, 1, 2, 4),
-        (4, 32, 32, 5, 32, 1, 2, 8),
-        (4, 32, 32, 5, 32, 1, 2, 8),
-        (4, 32, 32, 5, 32, 1, 2, 16),
-        (4, 32, 32, 5, 32, 1, 2, 32),
-        (4, 32, 256, 5, 512, 1, 2, 2),
-        (4, 32, 256, 5, 512, 1, 2, 128),
-        (4, 32, 256, 5, 512, 1, 2, 256),
-    )
-
-    for dtype in dtypes:
-        print("(N,  iH,  C),  (O,  wH,  C),   dtype,  stride, pads, groups, diff%")
-        for N, iH, C, wH, O, strides, padding, groups in shapes:
-            np_dtype = getattr(np, dtype)
-            time_mlx, time_torch = bench_shape(
-                N, iH, C, wH, O, strides, padding, np_dtype, groups
-            )
-            diff = time_torch / time_mlx - 1.0
-
-            print(
-                f"({N}, {iH:3d}, {C:3d}), ({O:3d}, {wH:2d}, {C:3d}), {dtype}, {strides:5d}, {padding:4d}, {groups:6d}, {100. * diff:+5.2f}%"
-            )
-
-            if time_mlx >= 2.0 * time_torch:
-                print("ATTENTION ^^^^^^^")

benchmarks/python/conv_bench.py

@@ -28,11 +28,11 @@ def bench(f, a, b):
     return (e - s) * 1e-9
 
 
-def make_mx_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
+def make_mx_conv_2D(strides=(1, 1), padding=(0, 0)):
     def mx_conv_2D(a, b):
         ys = []
         for i in range(N_iter_func):
-            y = mx.conv2d(a, b, stride=strides, padding=padding, groups=groups)
+            y = mx.conv2d(a, b, stride=strides, padding=padding)
             ys.append(y)
         mx.eval(ys)
         return ys
@@ -40,12 +40,12 @@ def make_mx_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
     return mx_conv_2D
 
 
-def make_pt_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
+def make_pt_conv_2D(strides=(1, 1), padding=(0, 0)):
     @torch.no_grad()
     def pt_conv_2D(a, b):
         ys = []
         for i in range(N_iter_func):
-            y = torch.conv2d(a, b, stride=strides, padding=padding, groups=groups)
+            y = torch.conv2d(a, b, stride=strides, padding=padding)
             ys.append(y)
         torch.mps.synchronize()
         return ys
@@ -53,12 +53,11 @@ def make_pt_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
     return pt_conv_2D
 
 
-def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype):
+def bench_shape(N, H, W, C, kH, kW, O, strides, padding, np_dtype):
+
     scale = 1.0 / math.sqrt(kH * kH * C)
     a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype)
-    b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype(
-        np_dtype
-    )
+    b_np = np.random.uniform(-scale, scale, (O, kH, kW, C)).astype(np_dtype)
 
     a_mx = mx.array(a_np)
     b_mx = mx.array(b_np)
@@ -68,15 +67,15 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype):
 
     torch.mps.synchronize()
 
-    f_mx = make_mx_conv_2D(strides, padding, groups)
-    f_pt = make_pt_conv_2D(strides, padding, groups)
+    f_mx = make_mx_conv_2D(strides, padding)
+    f_pt = make_pt_conv_2D(strides, padding)
 
     time_torch = bench(f_pt, a_pt, b_pt)
     time_mlx = bench(f_mx, a_mx, b_mx)
 
-    out_mx = mx.conv2d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    out_mx = mx.conv2d(a_mx, b_mx, stride=strides, padding=padding)
     out_pt = torch.conv2d(
-        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding
     )
     out_pt = torch.permute(out_pt, (0, 2, 3, 1))
     out_pt = out_pt.numpy(force=True)
@@ -85,7 +84,7 @@ def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype):
 
     if not np.allclose(out_pt, out_mx, atol=atol):
         print(
-            f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
+            f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
         )
 
     return time_mlx, time_torch
@@ -96,40 +95,35 @@ if __name__ == "__main__":
 
     dtypes = ("float32",)
     shapes = (
-        (4, 32, 32, 32, 5, 5, 32, (1, 1), (2, 2), 1),
-        (4, 32, 32, 64, 5, 5, 64, (1, 1), (2, 2), 1),
-        (4, 32, 32, 128, 5, 5, 128, (1, 1), (2, 2), 1),
-        (4, 32, 32, 256, 5, 5, 256, (1, 1), (2, 2), 1),
-        (4, 32, 32, 512, 5, 5, 512, (1, 1), (2, 2), 1),
-        (4, 64, 64, 32, 5, 5, 32, (1, 1), (2, 2), 1),
-        (4, 64, 64, 64, 5, 5, 64, (1, 1), (2, 2), 1),
-        (4, 64, 64, 128, 5, 5, 128, (1, 1), (2, 2), 1),
-        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2), 1),
-        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2), 2),
-        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2), 16),
-        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2), 64),
-        (4, 128, 128, 32, 5, 5, 32, (1, 1), (2, 2), 1),
-        (4, 128, 128, 64, 5, 5, 64, (1, 1), (2, 2), 1),
-        (4, 128, 128, 128, 5, 5, 128, (1, 1), (2, 2), 1),
-        (4, 256, 256, 32, 5, 5, 3, (1, 1), (2, 2), 1),
-        (4, 256, 256, 3, 5, 5, 32, (1, 1), (2, 2), 1),
-        (4, 128, 128, 64, 5, 5, 3, (1, 1), (2, 2), 1),
-        (4, 128, 128, 3, 5, 5, 64, (1, 1), (2, 2), 1),
+        (4, 32, 32, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 32, 32, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 32, 32, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 32, 32, 256, 5, 5, 256, (1, 1), (2, 2)),
+        (4, 32, 32, 512, 5, 5, 512, (1, 1), (2, 2)),
+        (4, 64, 64, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 64, 64, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 64, 64, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2)),
+        (4, 128, 128, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 128, 128, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 128, 128, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 256, 256, 32, 5, 5, 3, (1, 1), (2, 2)),
+        (4, 256, 256, 3, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 128, 128, 64, 5, 5, 3, (1, 1), (2, 2)),
+        (4, 128, 128, 3, 5, 5, 64, (1, 1), (2, 2)),
     )
 
     for dtype in dtypes:
-        print(
-            "(N,   H,   W,   C), (  O, kH, kW,   C),   dtype, stride,   pads,  groups, diff%"
-        )
-        for N, H, W, C, kH, kW, O, strides, padding, groups in shapes:
+        print("(N,   H,   W,   C), (  O, kH, kW,   C),   dtype, stride,   pads,  diff%")
+        for N, H, W, C, kH, kW, O, strides, padding in shapes:
             np_dtype = getattr(np, dtype)
             time_mlx, time_torch = bench_shape(
-                N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype
+                N, H, W, C, kH, kW, O, strides, padding, np_dtype
             )
             diff = time_torch / time_mlx - 1.0
 
             print(
-                f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%"
+                f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {100. * diff:+5.2f}%"
             )
             if time_mlx >= 2.0 * time_torch:
                 print("ATTENTION ^^^^^^^")

benchmarks/python/einsum_bench.py

@@ -1,84 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-import time
-
-import mlx.core as mx
-import numpy as np
-
-
-def timeit(fn, its=100, args=[]):
-    for _ in range(5):
-        fn(*args)
-    tic = time.perf_counter()
-    for _ in range(its):
-        fn(*args)
-    toc = time.perf_counter()
-    return 1e3 * (toc - tic) / its
-
-
-def time_little_einsum_path():
-    subscripts = "ik,kj->ij"
-    x = mx.ones((32, 32))
-    y = mx.ones((32, 32))
-    mx_time = timeit(mx.einsum_path, args=(subscripts, x, y))
-
-    x = np.array(x)
-    y = np.array(y)
-    np_time = timeit(np.einsum_path, args=(subscripts, x, y))
-    print("Timing little einsum path...")
-    print(f"MLX ... {mx_time:.3f} ms")
-    print(f"NumPy... {np_time:.3f} ms")
-
-
-def time_big_einsum_path():
-    chars = list("abcdefgh")
-    char_to_dim = {c: v for v, c in enumerate(chars)}
-
-    num_inputs = 10
-    inputs = []
-    subscripts = []
-    for _ in range(num_inputs):
-        subscript = np.random.choice(chars, size=5, replace=False).tolist()
-        subscripts.append("".join(subscript))
-        inputs.append(np.ones(list(char_to_dim[c] for c in subscript)))
-    subscripts = ",".join(subscripts)
-
-    np_time = timeit(np.einsum_path, args=(subscripts, *inputs))
-
-    inputs = [mx.array(x) for x in inputs]
-    mx_time = timeit(mx.einsum_path, args=(subscripts, *inputs))
-    print("Timing big einsum path...")
-    print(f"MLX ... {mx_time:.3f} ms")
-    print(f"NumPy... {np_time:.3f} ms")
-
-
-def time_attention():
-    def regular_attention(x):
-        # shape [batch, sequence, num_heads, head_dim]
-        queries, keys, values = x, x, x
-        scores = queries.transpose(0, 2, 1, 3) @ keys.transpose(0, 2, 3, 1)
-        scores = mx.softmax(scores, axis=-1)
-        output = (scores @ values.transpose(0, 2, 1, 3)).swapaxes(1, 2)
-        mx.eval(output)
-
-    def einsum_attention(x):
-        # shape [batch, sequence, num_heads, head_dim]
-        queries, keys, values = x, x, x
-        scores = mx.einsum("itjk,iujk->ijtu", queries, keys)
-        scores = mx.softmax(scores, axis=-1)
-        output = mx.einsum("ijtu,iujk->itjk", scores, values)
-        mx.eval(output)
-
-    x = mx.random.uniform(shape=(8, 512, 32, 128))
-
-    regular_time = timeit(regular_attention, args=(x,))
-    ein_time = timeit(einsum_attention, args=(x,))
-    print("Timing einsum attention...")
-    print(f"Regular ... {regular_time:.3f} ms")
-    print(f"Einsum ... {ein_time:.3f} ms")
-
-
-if __name__ == "__main__":
-    time_little_einsum_path()
-    time_big_einsum_path()
-    time_attention()

benchmarks/python/fft_bench.py

@@ -1,118 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-import matplotlib
-import mlx.core as mx
-import numpy as np
-import sympy
-import torch
-from time_utils import measure_runtime
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-
-def bandwidth_gb(runtime_ms, system_size):
-    bytes_per_fft = np.dtype(np.complex64).itemsize * 2
-    bytes_per_gb = 1e9
-    ms_per_s = 1e3
-    return system_size * bytes_per_fft / runtime_ms * ms_per_s / bytes_per_gb
-
-
-def run_bench(system_size, fft_sizes, backend="mlx", dim=1):
-    def fft_mlx(x):
-        if dim == 1:
-            out = mx.fft.fft(x)
-        elif dim == 2:
-            out = mx.fft.fft2(x)
-        mx.eval(out)
-        return out
-
-    def fft_mps(x):
-        if dim == 1:
-            out = torch.fft.fft(x)
-        elif dim == 2:
-            out = torch.fft.fft2(x)
-        torch.mps.synchronize()
-        return out
-
-    bandwidths = []
-    for n in fft_sizes:
-        batch_size = system_size // n**dim
-        shape = [batch_size] + [n for _ in range(dim)]
-        if backend == "mlx":
-            x_np = np.random.uniform(size=(system_size // n, n)).astype(np.complex64)
-            x = mx.array(x_np)
-            mx.eval(x)
-            fft = fft_mlx
-        elif backend == "mps":
-            x_np = np.random.uniform(size=(system_size // n, n)).astype(np.complex64)
-            x = torch.tensor(x_np, device="mps")
-            torch.mps.synchronize()
-            fft = fft_mps
-        else:
-            raise NotImplementedError()
-        runtime_ms = measure_runtime(fft, x=x)
-        bandwidth = bandwidth_gb(runtime_ms, np.prod(shape))
-        print(n, bandwidth)
-        bandwidths.append(bandwidth)
-
-    return np.array(bandwidths)
-
-
-def time_fft():
-    x = np.array(range(2, 512))
-    system_size = int(2**26)
-
-    print("MLX GPU")
-    with mx.stream(mx.gpu):
-        gpu_bandwidths = run_bench(system_size=system_size, fft_sizes=x)
-
-    print("MPS GPU")
-    mps_bandwidths = run_bench(system_size=system_size, fft_sizes=x, backend="mps")
-
-    print("CPU")
-    system_size = int(2**20)
-    with mx.stream(mx.cpu):
-        cpu_bandwidths = run_bench(system_size=system_size, fft_sizes=x)
-
-    x = np.array(x)
-
-    all_indices = x - x[0]
-    radix_2to13 = (
-        np.array([i for i in x if all(p <= 13 for p in sympy.primefactors(i))]) - x[0]
-    )
-    bluesteins = (
-        np.array([i for i in x if any(p > 13 for p in sympy.primefactors(i))]) - x[0]
-    )
-
-    for indices, name in [
-        (all_indices, "All"),
-        (radix_2to13, "Radix 2-13"),
-        (bluesteins, "Bluestein's"),
-    ]:
-        # plot bandwidths
-        print(name)
-        plt.scatter(x[indices], gpu_bandwidths[indices], color="green", label="GPU")
-        plt.scatter(x[indices], mps_bandwidths[indices], color="blue", label="MPS")
-        plt.scatter(x[indices], cpu_bandwidths[indices], color="red", label="CPU")
-        plt.title(f"MLX FFT Benchmark -- {name}")
-        plt.xlabel("N")
-        plt.ylabel("Bandwidth (GB/s)")
-        plt.legend()
-        plt.savefig(f"{name}.png")
-        plt.clf()
-
-    av_gpu_bandwidth = np.mean(gpu_bandwidths)
-    av_mps_bandwidth = np.mean(mps_bandwidths)
-    av_cpu_bandwidth = np.mean(cpu_bandwidths)
-    print("Average bandwidths:")
-    print("GPU:", av_gpu_bandwidth)
-    print("MPS:", av_mps_bandwidth)
-    print("CPU:", av_cpu_bandwidth)
-
-    portion_faster = len(np.where(gpu_bandwidths > mps_bandwidths)[0]) / len(x)
-    print("Percent MLX faster than MPS: ", portion_faster * 100)
-
-
-if __name__ == "__main__":
-    time_fft()

benchmarks/python/hadamard_bench.py

@@ -1,70 +0,0 @@
-import argparse
-
-import matplotlib
-import mlx.core as mx
-import numpy as np
-from time_utils import measure_runtime
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-
-
-def had(x):
-    y = mx.hadamard_transform(x)
-    mx.eval(y)
-
-
-def copy(x):
-    y = x + 1.0
-    mx.eval(y)
-
-
-def run(dtype):
-    system_size = 2**26
-    outputs = {}
-    for test_fn in (had, copy):
-        for m in [1, 12, 20, 28]:
-            if test_fn == copy:
-                key = "copy"
-            elif m == 1:
-                key = "had_2^k"
-            else:
-                key = "had_m*2^k"
-            outputs.setdefault(key, {})
-            for k in range(7, 14):
-                n = m * 2**k
-                if n > 2**15:
-                    continue
-                x_np = np.random.normal(size=(system_size // n, n)).astype(dtype)
-                x = mx.array(x_np)
-                runtime_ms = measure_runtime(test_fn, x=x)
-                bytes_per_gb = 1e9
-                ms_per_s = 1e3
-                bytes_per_had = np.dtype(x_np.dtype).itemsize * 2
-                bandwidth_gb = (
-                    system_size * bytes_per_had / runtime_ms * ms_per_s / bytes_per_gb
-                )
-                print(n, bandwidth_gb)
-                outputs[key][n] = bandwidth_gb
-
-    colors = {
-        "copy": "black",
-        "had_2^k": "steelblue",
-        "had_m*2^k": "skyblue",
-    }
-    for key, output in outputs.items():
-        plt.scatter(output.keys(), output.values(), color=colors[key], label=key)
-    plt.title(f"MLX Hadamard Benchmark -- {dtype.__name__}")
-    plt.xlabel("N")
-    plt.ylabel("Bandwidth (GB/s)")
-    plt.legend()
-    plt.savefig(f"bench_{dtype.__name__}.png")
-    plt.clf()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--fp16", action="store_true")
-    args = parser.parse_args()
-    dtype = np.float16 if args.fp16 else np.float32
-    run(dtype)

benchmarks/python/layer_norm_bench.py

@@ -1,41 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import mlx.core as mx
-import mlx.nn as nn
-from time_utils import time_fn
-
-
-def layer_norm(x, w, b, eps):
-    ot = x.dtype
-    x = x.astype(mx.float32)
-    mu = mx.mean(x, -1, keepdims=True)
-    v = mx.var(x, -1, keepdims=True)
-    return (x - mu) * mx.rsqrt(v + eps) * w + b
-
-
-def time_layer_norm():
-    f1 = lambda x, w, b, y: (layer_norm(x, w, b, 1e-5) * y).sum()
-    f2 = lambda x, w, b, y: (mx.fast.layer_norm(x, w, b, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0, 1, 2))
-    g2 = mx.grad(f2, argnums=(0, 1, 2))
-
-    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
-    b = mx.random.uniform(shape=(4096,)).astype(mx.float16)
-    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    mx.eval(x, w, b, y)
-
-    def layer_norm_loop(g, x, w, b):
-        gx, gw, gb = x, w, b
-        for _ in range(32):
-            gx, gw, gb = g(gx, gw, gb, y)
-        return gx, gw, gb
-
-    time_fn(layer_norm_loop, g1, x, w, b)
-    time_fn(layer_norm_loop, g2, x, w, b)
-    time_fn(layer_norm_loop, mx.compile(g1), x, w, b)
-    time_fn(layer_norm_loop, mx.compile(g2), x, w, b)
-
-
-if __name__ == "__main__":
-    time_layer_norm()

benchmarks/python/rms_norm_bench.py

@@ -1,39 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import mlx.core as mx
-import mlx.nn as nn
-from time_utils import time_fn
-
-
-def rms_norm(x, w, eps):
-    ot = x.dtype
-    x = x.astype(mx.float32)
-    n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
-    return (x * n).astype(ot) * w
-
-
-def time_rms_norm():
-    f1 = lambda x, w, y: (rms_norm(x, w, 1e-5) * y).sum()
-    f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0, 1))
-    g2 = mx.grad(f2, argnums=(0, 1))
-
-    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
-    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    mx.eval(x, w, y)
-
-    def rms_norm_loop(g, x, w):
-        gx, gw = x, w
-        for _ in range(32):
-            gx, gw = g(gx, gw, y)
-        return gx, gw
-
-    time_fn(rms_norm_loop, g1, x, w)
-    time_fn(rms_norm_loop, g2, x, w)
-    time_fn(rms_norm_loop, mx.compile(g1), x, w)
-    time_fn(rms_norm_loop, mx.compile(g2), x, w)
-
-
-if __name__ == "__main__":
-    time_rms_norm()

benchmarks/python/rope_bench.py

@@ -6,21 +6,21 @@ from time_utils import time_fn
 
 
 def time_rope():
-    rope = nn.RoPE(64)
+    rope = nn.RoPE(4096)
 
     # vec
-    x = mx.random.uniform(shape=(1, 32, 1, 128)).astype(mx.float16)
+    x = mx.random.uniform(shape=(1, 4096)).astype(mx.float16)
     mx.eval(x)
 
     def rope_vec(x):
         for _ in range(32):
-            x = rope(x, offset=100)
+            x = rope(x)
         return x
 
     time_fn(rope_vec, x)
 
     # matrix
-    x = mx.random.uniform(shape=(1, 32, 1024, 128)).astype(mx.float16)
+    x = mx.random.uniform(shape=(1024, 4096)).astype(mx.float16)
     mx.eval(x)
 
     def rope_mat(x):

benchmarks/python/sdpa_bench.py

@@ -1,62 +0,0 @@
-import argparse
-import math
-
-import mlx.core as mx
-from time_utils import time_fn
-
-MAX_SEQ = 300
-START_SEQ = 100
-SEQ_INCREMENT = 50
-
-
-def time_self_attention_primitives():
-    mx.random.seed(3)
-    B = 2
-    H = 38
-    D = 64
-    for R in range(START_SEQ, MAX_SEQ, SEQ_INCREMENT):
-        q = mx.random.uniform(shape=(B, H, R, D))
-        k = mx.random.uniform(shape=(B, H, R, D))
-        v = mx.random.uniform(shape=(B, H, R, D))
-        scale = 1.0 / math.sqrt(float(D))
-        mx.eval(q, k, v)
-
-        def sdpa_primitives(qs, ks, vs, alpha):
-            s = (alpha * qs) @ ks.transpose(0, 1, 3, 2)
-            p = mx.softmax(s.astype(mx.float32), axis=-1).astype(s.dtype)
-            o = p @ vs
-            return o
-
-        time_fn(sdpa_primitives, q, k, v, scale)
-
-
-def time_self_attention_sdpa():
-    mx.random.seed(3)
-    B = 2
-    H = 38
-    D = 64
-    for R in range(START_SEQ, MAX_SEQ, SEQ_INCREMENT):
-        q = mx.random.uniform(shape=(B, H, R, D))
-        k = mx.random.uniform(shape=(B, H, R, D))
-        v = mx.random.uniform(shape=(B, H, R, D))
-        scale = 1.0 / math.sqrt(float(D))
-        mx.eval(q, k, v)
-
-        def sdpa_fused(qs, ks, vs, alpha):
-            o = mx.fast.scaled_dot_product_attention(qs, ks, vs, scale=alpha)
-            return o
-
-        time_fn(sdpa_fused, q, k, v, scale)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser("MLX benchmarks.")
-    parser.add_argument("--gpu", action="store_true", help="Use the Metal back-end.")
-    args = parser.parse_args()
-    if args.gpu:
-        mx.set_default_device(mx.gpu)
-    else:
-        mx.set_default_device(mx.cpu)
-
-    time_self_attention_sdpa()
-    time_self_attention_primitives()

docs/Doxyfile

@@ -1,50 +0,0 @@
-################################################################################
-# Primary project setup.                                                       #
-################################################################################
-
-PROJECT_NAME           = "MLX"
-OUTPUT_DIRECTORY       = build
-XML_OUTPUT             = xml
-HTML_OUTPUT            = html
-STRIP_FROM_PATH        = ../
-INPUT                  = ../mlx
-FILE_PATTERNS          = *.h
-EXCLUDE_PATTERNS       = */private/*
-CREATE_SUBDIRS         = NO
-FULL_PATH_NAMES        = YES
-RECURSIVE              = YES
-GENERATE_HTML          = YES
-GENERATE_LATEX         = NO
-GENERATE_XML           = YES
-XML_PROGRAMLISTING     = YES
-
-################################################################################
-# Doxygen preprocessor / parser control.                                       #
-################################################################################
-
-ENABLE_PREPROCESSING   = YES
-MACRO_EXPANSION        = YES
-EXPAND_ONLY_PREDEF     = NO
-SKIP_FUNCTION_MACROS   = NO
-
-################################################################################
-# Compound extraction control.                                                 #
-################################################################################
-
-EXTRACT_ALL            = YES
-EXTRACT_PACKAGE        = YES
-EXTRACT_STATIC         = YES
-CASE_SENSE_NAMES       = NO
-
-################################################################################
-# Docstring control / customization.                                           #
-################################################################################
-
-JAVADOC_AUTOBRIEF      = YES
-
-################################################################################
-# Warning suppression.                                                         #
-################################################################################
-
-QUIET                  = YES
-WARN_IF_UNDOCUMENTED   = NO

docs/README.md

@@ -2,16 +2,12 @@
 
 ### Setup (do once)
 
-Install Doxygen:
+Install [sphinx](https://www.sphinx-doc.org/en/master/usage/installation.html)
+for example with `conda`:
 
 ```
-brew install doxygen
-```
-
-Install Python packages:
-
-```
-pip install -r requirements.txt
+conda install sphinx
+pip install sphinx-book-theme
 ```
 
 ### Build
@@ -19,7 +15,7 @@ pip install -r requirements.txt
 Build the docs from `mlx/docs/`
 
 ```
-doxygen && make html
+make html
 ```
 
 View the docs by running a server in `mlx/docs/build/html/`:

docs/requirements.txt

@@ -1,4 +0,0 @@
-sphinx
-breathe
-sphinx-book-theme
-mlx

docs/src/_templates/nn-module-template.rst

@@ -1,20 +0,0 @@
-{{ fullname | escape | underline}}
-
-.. currentmodule:: {{ module }}
-
-.. autoclass:: {{ objname }}
-
-   {% block methods %}
-
-   {% if methods %}
-   .. rubric:: {{ _('Methods') }}
-
-   .. autosummary::
-   {% for item in methods %}
-      {%- if item not in inherited_members and item != "__init__" %}
-         ~{{ name }}.{{ item }}
-      {%- endif %}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
-

docs/src/conf.py

@@ -22,7 +22,6 @@ extensions = [
     "sphinx.ext.autosummary",
     "sphinx.ext.intersphinx",
     "sphinx.ext.napoleon",
-    "breathe",
 ]
 
 python_use_unqualified_type_names = True
@@ -30,20 +29,16 @@ autosummary_generate = True
 autosummary_filename_map = {"mlx.core.Stream": "stream_class"}
 
 intersphinx_mapping = {
-    "python": ("https://docs.python.org/3", None),
-    "numpy": ("https://numpy.org/doc/stable/", None),
+    "https://docs.python.org/3": None,
+    "https://numpy.org/doc/stable/": None,
 }
 
-breathe_projects = {"mlx": "../build/xml"}
-breathe_default_project = "mlx"
-
 templates_path = ["_templates"]
 html_static_path = ["_static"]
 source_suffix = ".rst"
-main_doc = "index"
+master_doc = "index"
 highlight_language = "python"
 pygments_style = "sphinx"
-add_module_names = False
 
 # -- Options for HTML output -------------------------------------------------
 
@@ -64,34 +59,3 @@ html_theme_options = {
 # -- Options for HTMLHelp output ---------------------------------------------
 
 htmlhelp_basename = "mlx_doc"
-
-
-def setup(app):
-    from sphinx.util import inspect
-
-    wrapped_isfunc = inspect.isfunction
-
-    def isfunc(obj):
-        type_name = str(type(obj))
-        if "nanobind.nb_method" in type_name or "nanobind.nb_func" in type_name:
-            return True
-        return wrapped_isfunc(obj)
-
-    inspect.isfunction = isfunc
-
-
-# -- Options for LaTeX output ------------------------------------------------
-
-latex_documents = [(main_doc, "MLX.tex", "MLX Documentation", author, "manual")]
-latex_elements = {
-    "preamble": r"""
-    \usepackage{enumitem}
-    \setlistdepth{5}
-    \setlist[itemize,1]{label=$\bullet$}
-    \setlist[itemize,2]{label=$\bullet$}
-    \setlist[itemize,3]{label=$\bullet$}
-    \setlist[itemize,4]{label=$\bullet$}
-    \setlist[itemize,5]{label=$\bullet$}
-    \renewlist{itemize}{itemize}{5}
-""",
-}

docs/src/cpp/ops.rst

@@ -3,5 +3,4 @@
 Operations
 ==========
 
-.. doxygengroup:: ops
-   :content-only:
+

docs/src/dev/extensions.rst

@@ -1,16 +1,24 @@
-Custom Extensions in MLX
-========================
+Developer Documentation
+=======================
 
-You can extend MLX with custom operations on the CPU or GPU. This guide
-explains how to do that with a simple example.
+MLX provides a open and flexible backend to which users may add operations 
+and specialized implementations without much hassle. While the library supplies
+efficient operations that can be used and composed for any number of 
+applications, there may arise cases where new functionalities or highly 
+optimized implementations are needed. For such cases, you may design and 
+implement your own operations that link to and build on top of :mod:`mlx.core`.
+We will introduce the inner-workings of MLX and go over a simple example to 
+learn the steps involved in adding new operations to MLX with your own CPU 
+and GPU implementations. 
 
 Introducing the Example
 -----------------------
 
-Let's say you would like an operation that takes in two arrays, ``x`` and
-``y``, scales them both by coefficients ``alpha`` and ``beta`` respectively,
-and then adds them together to get the result ``z = alpha * x + beta * y``.
-You can do that in MLX directly:
+Let's say that you would like an operation that takes in two arrays, 
+``x`` and ``y``, scales them both by some coefficients ``alpha`` and ``beta``
+respectively, and then adds them together to get the result 
+``z = alpha * x + beta * y``. Well, you can very easily do that by just 
+writing out a function as follows:
 
 .. code-block:: python
 
@@ -19,35 +27,44 @@ You can do that in MLX directly:
     def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
         return alpha * x + beta * y
 
-This function performs that operation while leaving the implementation and
-function transformations to MLX.
+This function performs that operation while leaving the implementations and 
+differentiation to MLX. 
 
-However you may need to customize the underlying implementation, perhaps to
-make it faster or for custom differentiation. In this tutorial we will go
-through adding custom extensions. It will cover:
+However, you work with vector math libraries often and realize that the 
+``axpby`` routine defines the same operation ``Y = (alpha * X) + (beta * Y)``. 
+You would really like the part of your applications that does this operation 
+on the CPU to be very fast - so you decide that you want it to rely on the 
+``axpby`` routine provided by the Accelerate_ framework. Continuing to impose 
+our assumptions on to you, let's also assume that you want to learn how to add 
+your own implementation for the gradients of your new operation while going 
+over the ins-and-outs of the MLX framework. 
 
-* The structure of the MLX library.
-* Implementing a CPU operation that redirects to Accelerate_ when appropriate.
-* Implementing a GPU operation using metal.
-* Adding the ``vjp`` and ``jvp`` function transformation.
-* Building a custom extension and binding it to python.
+Well, what a coincidence! You are in the right place. Over the course of this 
+example, we will learn:
+
+* The structure of the MLX library from the frontend API to the backend implementations.
+* How to implement your own CPU backend that redirects to Accelerate_ when appropriate (and a fallback if needed).
+* How to implement your own GPU implementation using metal.
+* How to add your own ``vjp`` and ``jvp``.
+* How to build your implementations, link them to MLX, and bind them to python.
 
 Operations and Primitives
 -------------------------
 
-Operations in MLX build the computation graph. Primitives provide the rules for
-evaluating and transforming the graph. Let's start by discussing operations in
-more detail.
+In one sentence, operations in MLX build the computation graph, and primitives 
+provide the rules for evaluation and transformations of said graph. Let's start 
+by discussing operations in more detail. 
 
 Operations
 ^^^^^^^^^^^
 
-Operations are the front-end functions that operate on arrays. They are defined
-in the C++ API (:ref:`cpp_ops`), and the Python API (:ref:`ops`) binds them.
+Operations are the frontend functions that operate on arrays. They are defined 
+in the C++ API (:ref:`cpp_ops`) and then we provide bindings to these 
+operations in the Python API (:ref:`ops`). 
 
-We would like an operation, :meth:`axpby` that takes in two arrays ``x`` and
-``y``, and two scalars, ``alpha`` and ``beta``. This is how to define it in
-C++:
+We would like an operation, :meth:`axpby` that takes in two arrays ``x`` and ``y``,
+and two scalars, ``alpha`` and ``beta``. This is how we would define it in the 
+C++ API:
 
 .. code-block:: C++
 
@@ -66,7 +83,10 @@ C++:
         StreamOrDevice s = {} // Stream on which to schedule the operation
     );
 
-The simplest way to this operation is in terms of existing operations:
+
+This operation itself can call other operations within it if needed. So, the 
+simplest way to go about implementing this operation would be do so in terms 
+of existing operations. 
 
 .. code-block:: C++
 
@@ -80,23 +100,25 @@ The simplest way to this operation is in terms of existing operations:
         // Scale x and y on the provided stream
         auto ax = multiply(array(alpha), x, s);
         auto by = multiply(array(beta), y, s);
-
+        
         // Add and return
         return add(ax, by, s);
     }
 
-The operations themselves do not contain the implementations that act on the
-data, nor do they contain the rules of transformations. Rather, they are an
-easy to use interface that use :class:`Primitive` building blocks.
+However, as we discussed earlier, this is not our goal. The operations themselves 
+do not contain the implementations that act on the data, nor do they contain the
+rules of transformations. Rather, they are an easy to use interface that build 
+on top of the building blocks we call :class:`Primitive`. 
 
 Primitives
 ^^^^^^^^^^^
 
-A :class:`Primitive` is part of the computation graph of an :class:`array`. It
-defines how to create outputs arrays given a input arrays. Further, a
-:class:`Primitive` has methods to run on the CPU or GPU and for function
-transformations such as ``vjp`` and ``jvp``.  Lets go back to our example to be
-more concrete:
+A :class:`Primitive` is part of the computation graph of an :class:`array`. It 
+defines how to create an output given a set of input :class:`array` . Further,
+a :class:`Primitive` is a class that contains rules on how it is evaluated 
+on the CPU or GPU, and how it acts under transformations such as ``vjp`` and 
+``jvp``. These words on their own can be a bit abstract, so lets take a step 
+back and go to our example to give ourselves a more concrete image. 
 
 .. code-block:: C++
 
@@ -112,15 +134,11 @@ more concrete:
         * To avoid unnecessary allocations, the evaluation function
         * is responsible for allocating space for the array.
         */
-        void eval_cpu(
-            const std::vector<array>& inputs,
-            std::vector<array>& outputs) override;
-        void eval_gpu(
-            const std::vector<array>& inputs,
-            std::vector<array>& outputs) override;
+        void eval_cpu(const std::vector<array>& inputs, array& out) override;
+        void eval_gpu(const std::vector<array>& inputs, array& out) override;
 
         /** The Jacobian-vector product. */
-        std::vector<array> jvp(
+        array jvp(
             const std::vector<array>& primals,
             const std::vector<array>& tangents,
             const std::vector<int>& argnums) override;
@@ -129,8 +147,7 @@ more concrete:
         std::vector<array> vjp(
             const std::vector<array>& primals,
             const array& cotan,
-            const std::vector<int>& argnums,
-            const std::vector<array>& outputs) override;
+            const std::vector<int>& argnums) override;
 
         /**
         * The primitive must know how to vectorize itself across
@@ -138,7 +155,7 @@ more concrete:
         * representing the vectorized computation and the axis which
         * corresponds to the output vectorized dimension.
         */
-        virtual std::pair<std::vector<array>, std::vector<int>> vmap(
+        std::pair<array, int> vmap(
             const std::vector<array>& inputs,
             const std::vector<int>& axes) override;
 
@@ -158,22 +175,22 @@ more concrete:
         void eval(const std::vector<array>& inputs, array& out);
     };
 
-The :class:`Axpby` class derives from the base :class:`Primitive` class. The
-:class:`Axpby` treats ``alpha`` and ``beta`` as parameters. It then provides
-implementations of how the output array is produced given the inputs through
-:meth:`Axpby::eval_cpu` and :meth:`Axpby::eval_gpu`. It also provides rules
-of transformations in :meth:`Axpby::jvp`, :meth:`Axpby::vjp`, and
-:meth:`Axpby::vmap`.
+The :class:`Axpby` class derives from the base :class:`Primitive` class and 
+follows the above demonstrated interface. :class:`Axpby` treats ``alpha`` and 
+``beta`` as parameters. It then provides implementations of how the array ``out`` 
+is produced given ``inputs`` through :meth:`Axpby::eval_cpu` and 
+:meth:`Axpby::eval_gpu`. Further, it provides rules of transformations in 
+:meth:`Axpby::jvp`, :meth:`Axpby::vjp`, and :meth:`Axpby::vmap`. 
 
-Using the Primitive
-^^^^^^^^^^^^^^^^^^^
+Using the Primitives
+^^^^^^^^^^^^^^^^^^^^^
 
-Operations can use this :class:`Primitive` to add a new :class:`array` to the
-computation graph. An :class:`array` can be constructed by providing its data
-type, shape, the :class:`Primitive` that computes it, and the :class:`array`
-inputs that are passed to the primitive.
+Operations can use this :class:`Primitive` to add a new :class:`array` to 
+the computation graph. An :class:`array` can be constructed by providing its 
+data type, shape, the :class:`Primitive` that computes it, and the 
+:class:`array` inputs that are passed to the primitive.
 
-Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
+Let's re-implement our operation now in terms of our :class:`Axpby` primitive.
 
 .. code-block:: C++
 
@@ -206,7 +223,7 @@ Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
             /* const std::vector<int>& shape = */ out_shape,
             /* Dtype dtype = */ out_dtype,
             /* std::unique_ptr<Primitive> primitive = */
-            std::make_shared<Axpby>(to_stream(s), alpha, beta),
+            std::make_unique<Axpby>(to_stream(s), alpha, beta),
             /* const std::vector<array>& inputs = */ broadcasted_inputs);
     }
 
@@ -221,26 +238,27 @@ This operation now handles the following:
 Implementing the Primitive
 --------------------------
 
-No computation happens when we call the operation alone. The operation only
-builds the computation graph. When we evaluate the output array, MLX schedules
-the execution of the computation graph, and calls :meth:`Axpby::eval_cpu` or
-:meth:`Axpby::eval_gpu` depending on the stream/device specified by the user.
+No computation happens when we call the operation alone. In effect, the 
+operation only builds the computation graph. When we evaluate the output 
+array, MLX schedules the execution of the computation graph, and calls
+:meth:`Axpby::eval_cpu` or :meth:`Axpby::eval_gpu` depending on the 
+stream/device specified by the user. 
 
 .. warning::
     When :meth:`Primitive::eval_cpu` or :meth:`Primitive::eval_gpu` are called,
     no memory has been allocated for the output array. It falls on the implementation
-    of these functions to allocate memory as needed.
+    of these functions to allocate memory as needed
 
-Implementing the CPU Back-end
+Implementing the CPU Backend
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Let's start by implementing a naive and generic version of
-:meth:`Axpby::eval_cpu`. We declared this as a private member function of
-:class:`Axpby` earlier called :meth:`Axpby::eval`.
+Let's start by trying to implement a naive and generic version of 
+:meth:`Axpby::eval_cpu`. We declared this as a private member function of 
+:class:`Axpby` earlier called :meth:`Axpby::eval`. 
 
-Our naive method will go over each element of the output array, find the
-corresponding input elements of ``x`` and ``y`` and perform the operation
-point-wise. This is captured in the templated function :meth:`axpby_impl`.
+Our naive method will go over each element of the output array, find the 
+corresponding input elements of ``x`` and ``y`` and perform the operation 
+pointwise. This is captured in the templated function :meth:`axpby_impl`. 
 
 .. code-block:: C++
 
@@ -278,19 +296,19 @@ point-wise. This is captured in the templated function :meth:`axpby_impl`.
         }
     }
 
-Our implementation should work for all incoming floating point arrays.
-Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
-``complex64``. We throw an error if we encounter an unexpected type.
+Now, we would like our implementation to be able to do this pointwise operation 
+for all incoming floating point arrays. Accordingly, we add dispatches for 
+``float32``, ``float16``, ``bfloat16`` and ``complex64``. We throw an error 
+if we encounter an unexpected type.
 
 .. code-block:: C++
 
     /** Fall back implementation for evaluation on CPU */
-    void Axpby::eval(
-      const std::vector<array>& inputs,
-      const std::vector<array>& outputs) {
+    void Axpby::eval(const std::vector<array>& inputs, array& out) {
+        // Check the inputs (registered in the op while constructing the out array)
+        assert(inputs.size() == 2);
         auto& x = inputs[0];
         auto& y = inputs[1];
-        auto& out = outputs[0];
 
         // Dispatch to the correct dtype
         if (out.dtype() == float32) {
@@ -303,26 +321,28 @@ Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
             return axpby_impl<complex64_t>(x, y, out, alpha_, beta_);
         } else {
             throw std::runtime_error(
-                "[Axpby] Only supports floating point types.");
+                "Axpby is only supported for floating point types.");
         }
     }
 
-This is good as a fallback implementation. We can use the ``axpby`` routine
-provided by the Accelerate_ framework for a faster implementation in certain
-cases:
+We have a fallback implementation! Now, to do what we are really here to do. 
+Remember we wanted to use the ``axpby`` routine provided by the Accelerate_
+framework? Well, there are 3 complications to keep in mind:
 
 #.  Accelerate does not provide implementations of ``axpby`` for half precision
-    floats. We can only use it for ``float32`` types.
-#.  Accelerate assumes the inputs ``x`` and ``y`` are contiguous and all
-    elements have fixed strides between them. We only direct to Accelerate
-    if both ``x`` and ``y`` are row contiguous or column contiguous.
-#.  Accelerate performs the routine ``Y = (alpha * X) + (beta * Y)`` in-place.
-    MLX expects to write the output to a new array. We must copy the elements
-    of ``y`` into the output and use that as an input to ``axpby``.
+    floats. We can only direct to it for ``float32`` types 
+#.  Accelerate assumes the inputs ``x`` and ``y`` are contiguous and all elements
+    have fixed strides between them. Possibly due to broadcasts and transposes, 
+    we aren't guaranteed that the inputs fit this requirement. We can 
+    only direct to Accelerate if both ``x`` and ``y`` are row contiguous or 
+    column contiguous. 
+#.  Accelerate performs the routine ``Y = (alpha * X) + (beta * Y)`` inplace. 
+    MLX expects to write out the answer to a new array. We must copy the elements 
+    of ``y`` into the output array and use that as an input to ``axpby``
 
-Let's write an implementation that uses Accelerate in the right conditions.
-It allocates data for the output, copies ``y`` into it, and then calls the
-:func:`catlas_saxpby` from accelerate.
+Let's write out an implementation that uses Accelerate in the right conditions. 
+It must simply allocate data for the output, copy elements of ``y`` into it, 
+and then call the :meth:`catlas_saxpby` from accelerate. 
 
 .. code-block:: C++
 
@@ -336,7 +356,17 @@ It allocates data for the output, copies ``y`` into it, and then calls the
         // Accelerate library provides catlas_saxpby which does
         // Y = (alpha * X) + (beta * Y) in place
         // To use it, we first copy the data in y over to the output array
-        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+
+        // This specialization requires both x and y be contiguous in the same mode
+        // i.e: corresponding linear indices in both point to corresponding elements
+        // The data in the output array is allocated to match the strides in y
+        // such that x, y, and out are contiguous in the same mode and
+        // no transposition is needed
+        out.set_data(
+            allocator::malloc_or_wait(y.data_size() * out.itemsize()),
+            y.data_size(),
+            y.strides(),
+            y.flags());
 
         // We then copy over the elements using the contiguous vector specialization
         copy_inplace(y, out, CopyType::Vector);
@@ -359,20 +389,18 @@ It allocates data for the output, copies ``y`` into it, and then calls the
             /* INCY = */ 1);
     }
 
-For inputs that do not fit the criteria for accelerate, we fall back to
-:meth:`Axpby::eval`. With this in mind, let's finish our
-:meth:`Axpby::eval_cpu`.
+Great! But what about the inputs that do not fit the criteria for accelerate?
+Luckily, we can always just direct back to :meth:`Axpby::eval`.
+
+With this in mind, lets finally implement our :meth:`Axpby::eval_cpu`.
 
 .. code-block:: C++
 
     /** Evaluate primitive on CPU using accelerate specializations */
-    void Axpby::eval_cpu(
-      const std::vector<array>& inputs,
-      const std::vector<array>& outputs) {
+    void Axpby::eval_cpu(const std::vector<array>& inputs, array& out) {
         assert(inputs.size() == 2);
         auto& x = inputs[0];
         auto& y = inputs[1];
-        auto& out = outputs[0];
 
         // Accelerate specialization for contiguous single precision float arrays
         if (out.dtype() == float32 &&
@@ -382,33 +410,35 @@ For inputs that do not fit the criteria for accelerate, we fall back to
             return;
         }
 
-        // Fall back to common back-end if specializations are not available
-        eval(inputs, outputs);
+        // Fall back to common backend if specializations are not available
+        eval(inputs, out);
     }
 
-Just this much is enough to run the operation :meth:`axpby` on a CPU stream! If
-you do not plan on running the operation on the GPU or using transforms on
-computation graphs that contain :class:`Axpby`, you can stop implementing the
-primitive here and enjoy the speed-ups you get from the Accelerate library.
+We have now hit a milestone! Just this much is enough to run the operation 
+:meth:`axpby` on a CPU stream! 
 
-Implementing the GPU Back-end
+If you do not plan on running the operation on the GPU or using transforms on 
+computation graphs that contain :class:`Axpby`, you can stop implementing the 
+primitive here and enjoy the speed-ups you get from the Accelerate library. 
+
+Implementing the GPU Backend
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Apple silicon devices address their GPUs using the Metal_ shading language, and
-GPU kernels in MLX are written using Metal.
+Apple silicon devices address their GPUs using the Metal_ shading language, and 
+all GPU kernels in MLX are written using metal. 
 
 .. note::
 
-    Here are some helpful resources if you are new to Metal:
+    Here are some helpful resources if you are new to metal!
 
     * A walkthrough of the metal compute pipeline: `Metal Example`_
     * Documentation for metal shading language: `Metal Specification`_
     * Using metal from C++: `Metal-cpp`_
 
-Let's keep the GPU kernel simple. We will launch exactly as many threads as
-there are elements in the output. Each thread will pick the element it needs
-from ``x`` and ``y``, do the point-wise operation, and update its assigned
-element in the output.
+Let's keep the GPU algorithm simple. We will launch exactly as many threads 
+as there are elements in the output. Each thread will pick the element it needs 
+from ``x`` and ``y``, do the pointwise operation, and then update its assigned 
+element in the output. 
 
 .. code-block:: C++
 
@@ -427,14 +457,15 @@ element in the output.
         // Convert linear indices to offsets in array
         auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
         auto y_offset = elem_to_loc(index, shape, y_strides, ndim);
-
+        
         // Do the operation and update the output
-        out[index] =
+        out[index] = 
             static_cast<T>(alpha) * x[x_offset] + static_cast<T>(beta) * y[y_offset];
     }
 
 We then need to instantiate this template for all floating point types and give
-each instantiation a unique host name so we can identify it.
+each instantiation a unique host name so we can identify the right kernel for 
+each data type. 
 
 .. code-block:: C++
 
@@ -457,21 +488,29 @@ each instantiation a unique host name so we can identify it.
     instantiate_axpby(bfloat16, bfloat16_t);
     instantiate_axpby(complex64, complex64_t);
 
-The logic to determine the kernel, set the inputs, resolve the grid dimensions,
-and dispatch to the GPU are contained in :meth:`Axpby::eval_gpu` as shown
+This kernel will be compiled into a metal library ``mlx_ext.metallib`` as we 
+will see later in :ref:`Building with CMake`. In the following example, we 
+assume that the library ``mlx_ext.metallib`` will always be co-located with 
+the executable/ shared-library calling the :meth:`register_library` function. 
+The :meth:`register_library` function takes the library's name and potential 
+path (or in this case, a function that can produce the path of the metal 
+library) and tries to load that library if it hasn't already been registered 
+by the relevant static :class:`mlx::core::metal::Device` object. This is why, 
+it is important to package your C++ library with the metal library. We will 
+go over this process in more detail later. 
+
+The logic to determine the kernel, set the inputs, resolve the grid dimensions 
+and dispatch it to the GPU are contained in :meth:`Axpby::eval_gpu` as shown 
 below.
 
 .. code-block:: C++
 
     /** Evaluate primitive on GPU */
-    void Axpby::eval_gpu(
-      const std::vector<array>& inputs,
-      std::vector<array>& outputs) {
+    void Axpby::eval_gpu(const std::vector<array>& inputs, array& out) {
         // Prepare inputs
         assert(inputs.size() == 2);
         auto& x = inputs[0];
         auto& y = inputs[1];
-        auto& out = outputs[0];
 
         // Each primitive carries the stream it should execute on
         // and each stream carries its device identifiers
@@ -479,21 +518,22 @@ below.
         // We get the needed metal device using the stream
         auto& d = metal::device(s.device);
 
-        // Allocate output memory
+        // Allocate output memory 
         out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
-        // Resolve name of kernel
+        // Resolve name of kernel (corresponds to axpby.metal)
         std::ostringstream kname;
         kname << "axpby_" << "general_" << type_to_name(out);
 
-        // Make sure the metal library is available
-        d.register_library("mlx_ext");
+        // Make sure the metal library is available and look for it
+        // in the same folder as this executable if needed
+        d.register_library("mlx_ext", metal::get_colocated_mtllib_path);
 
         // Make a kernel from this metal library
         auto kernel = d.get_kernel(kname.str(), "mlx_ext");
 
         // Prepare to encode kernel
-        auto& compute_encoder = d.get_command_encoder(s.index);
+        auto compute_encoder = d.get_command_encoder(s.index);
         compute_encoder->setComputePipelineState(kernel);
 
         // Kernel parameters are registered with buffer indices corresponding to
@@ -502,17 +542,17 @@ below.
         size_t nelem = out.size();
 
         // Encode input arrays to kernel
-        compute_encoder.set_input_array(x, 0);
-        compute_encoder.set_input_array(y, 1);
+        set_array_buffer(compute_encoder, x, 0);
+        set_array_buffer(compute_encoder, y, 1);
 
         // Encode output arrays to kernel
-        compute_encoder.set_output_array(out, 2);
+        set_array_buffer(compute_encoder, out, 2);
 
         // Encode alpha and beta
         compute_encoder->setBytes(&alpha_, sizeof(float), 3);
         compute_encoder->setBytes(&beta_, sizeof(float), 4);
 
-        // Encode shape, strides and ndim
+        // Encode shape, strides and ndim 
         compute_encoder->setBytes(x.shape().data(), ndim * sizeof(int), 5);
         compute_encoder->setBytes(x.strides().data(), ndim * sizeof(size_t), 6);
         compute_encoder->setBytes(y.strides().data(), ndim * sizeof(size_t), 7);
@@ -530,30 +570,33 @@ below.
 
         // Launch the grid with the given number of threads divided among
         // the given threadgroups
-        compute_encoder.dispatchThreads(grid_dims, group_dims);
+        compute_encoder->dispatchThreads(grid_dims, group_dims);
     }
 
 We can now call the :meth:`axpby` operation on both the CPU and the GPU!
 
-A few things to note about MLX and Metal before moving on. MLX keeps track of
-the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is
-associated. We rely on :meth:`d.get_command_encoder` to give us the active
-metal compute command encoder instead of building a new one and calling
-:meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute
-pipelines) to the active command buffer until some specified limit is hit or
-the command buffer needs to be flushed for synchronization.
+A few things to note about MLX and metal before moving on. MLX keeps track 
+of the active ``compute_encoder``. We rely on :meth:`d.get_command_encoder` 
+to give us the active metal compute command encoder instead of building a 
+new one and calling :meth:`compute_encoder->end_encoding` at the end. 
+MLX keeps adding kernels (compute pipelines) to the active command encoder 
+until some specified limit is hit or the compute encoder needs to be flushed 
+for synchronization. MLX also handles enqueuing and committing the associated 
+command buffers as needed. We suggest taking a deeper dive into 
+:class:`metal::Device` if you would like to study this routine further.
 
 Primitive Transforms
 ^^^^^^^^^^^^^^^^^^^^^
 
-Next, let's add implementations for transformations in a :class:`Primitive`.
-These transformations can be built on top of other operations, including the
-one we just defined:
+Now that we have come this far, let's also learn how to add implementations to 
+transformations in a :class:`Primitive`. These transformations can be built on 
+top of our operations, including the one we just defined now. Which then gives 
+us the following :meth:`Axpby::jvp` and :meth:`Axpby::vjp` implementations.
 
 .. code-block:: C++
 
     /** The Jacobian-vector product. */
-    std::vector<array> Axpby::jvp(
+    array Axpby::jvp(
             const std::vector<array>& primals,
             const std::vector<array>& tangents,
             const std::vector<int>& argnums) {
@@ -568,12 +611,12 @@ one we just defined:
         if (argnums.size() > 1) {
             auto scale = argnums[0] == 0 ? alpha_ : beta_;
             auto scale_arr = array(scale, tangents[0].dtype());
-            return {multiply(scale_arr, tangents[0], stream())};
+            return multiply(scale_arr, tangents[0], stream());
         }
         // If, argnums = {0, 1}, we take contributions from both
         // which gives us jvp = tangent_x * alpha + tangent_y * beta
         else {
-            return {axpby(tangents[0], tangents[1], alpha_, beta_, stream())};
+            return axpby(tangents[0], tangents[1], alpha_, beta_, stream());
         }
     }
 
@@ -582,35 +625,34 @@ one we just defined:
     /** The vector-Jacobian product. */
     std::vector<array> Axpby::vjp(
             const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
-            const std::vector<int>& argnums,
-            const std::vector<int>& /* unused */) {
+            const array& cotan,
+            const std::vector<int>& argnums) {
         // Reverse mode diff
         std::vector<array> vjps;
         for (auto arg : argnums) {
             auto scale = arg == 0 ? alpha_ : beta_;
-            auto scale_arr = array(scale, cotangents[0].dtype());
-            vjps.push_back(multiply(scale_arr, cotangents[0], stream()));
+            auto scale_arr = array(scale, cotan.dtype());
+            vjps.push_back(multiply(scale_arr, cotan, stream()));
         }
         return vjps;
     }
 
-Note, a transformation does not need to be fully defined to start using
-the :class:`Primitive`.
+Finally, you need not have a transformation fully defined to start using your 
+own :class:`Primitive`.
 
 .. code-block:: C++
 
     /** Vectorize primitive along given axis */
-    std::pair<std::vector<array>, std::vector<int>> Axpby::vmap(
+    std::pair<array, int> Axpby::vmap(
             const std::vector<array>& inputs,
             const std::vector<int>& axes) {
-        throw std::runtime_error("[Axpby] vmap not implemented.");
+        throw std::runtime_error("Axpby has no vmap implementation.");
     }
 
 Building and Binding
 --------------------
 
-Let's look at the overall directory structure first.
+Let's look at the overall directory structure first. 
 
 | extensions
 | ├── axpby
@@ -624,39 +666,40 @@ Let's look at the overall directory structure first.
 | └── setup.py
 
 * ``extensions/axpby/`` defines the C++ extension library
-* ``extensions/mlx_sample_extensions`` sets out the structure for the
-  associated Python package
-* ``extensions/bindings.cpp`` provides Python bindings for our operation
-* ``extensions/CMakeLists.txt`` holds CMake rules to build the library and
-  Python bindings
+* ``extensions/mlx_sample_extensions`` sets out the structure for the 
+  associated python package
+* ``extensions/bindings.cpp`` provides python bindings for our operation
+* ``extensions/CMakeLists.txt`` holds CMake rules to build the library and 
+  python bindings
 * ``extensions/setup.py`` holds the ``setuptools`` rules to build and install
-  the Python package
+  the python package
 
 Binding to Python
 ^^^^^^^^^^^^^^^^^^
 
-We use nanobind_ to build a Python API for the C++ library. Since bindings for
+We use PyBind11_ to build a Python API for the C++ library. Since bindings for
 components such as :class:`mlx.core.array`, :class:`mlx.core.stream`, etc. are
-already provided, adding our :meth:`axpby` is simple.
+already provided, adding our :meth:`axpby` is simple!
 
 .. code-block:: C++
 
-   NB_MODULE(_ext, m) {
-        m.doc() = "Sample extension for MLX";
+    PYBIND11_MODULE(mlx_sample_extensions, m) {
+        m.doc() = "Sample C++ and metal extensions for MLX";
 
         m.def(
             "axpby",
             &axpby,
             "x"_a,
             "y"_a,
+            py::pos_only(),
             "alpha"_a,
             "beta"_a,
-            nb::kw_only(),
-            "stream"_a = nb::none(),
-            R"(
+            py::kw_only(),
+            "stream"_a = py::none(),
+            R"pbdoc(
                 Scale and sum two vectors element-wise
                 ``z = alpha * x + beta * y``
-
+                
                 Follows numpy style broadcasting between ``x`` and ``y``
                 Inputs are upcasted to floats if needed
 
@@ -668,17 +711,17 @@ already provided, adding our :meth:`axpby` is simple.
 
                 Returns:
                     array: ``alpha * x + beta * y``
-            )");
+            )pbdoc");
     }
 
-Most of the complexity in the above example comes from additional bells and
+Most of the complexity in the above example comes from additional bells and 
 whistles such as the literal names and doc-strings.
 
 .. warning::
 
-    :mod:`mlx.core` must be imported before importing
-    :mod:`mlx_sample_extensions` as defined by the nanobind module above to
-    ensure that the casters for :mod:`mlx.core` components like
+    :mod:`mlx.core` needs to be imported before importing 
+    :mod:`mlx_sample_extensions` as defined by the pybind11 module above to 
+    ensure that the casters for :mod:`mlx.core` components like 
     :class:`mlx.core.array` are available.
 
 .. _Building with CMake:
@@ -686,8 +729,8 @@ whistles such as the literal names and doc-strings.
 Building with CMake
 ^^^^^^^^^^^^^^^^^^^^
 
-Building the C++ extension library only requires that you ``find_package(MLX
-CONFIG)`` and then link it to your library.
+Building the C++ extension library itself is simple, it only requires that you 
+``find_package(MLX CONFIG)`` and then link it to your library. 
 
 .. code-block:: cmake
 
@@ -709,12 +752,12 @@ CONFIG)`` and then link it to your library.
     # Link to mlx
     target_link_libraries(mlx_ext PUBLIC mlx)
 
-We also need to build the attached Metal library. For convenience, we provide a
-:meth:`mlx_build_metallib` function that builds a ``.metallib`` target given
-sources, headers, destinations, etc. (defined in ``cmake/extension.cmake`` and
-automatically imported with MLX package).
+We also need to build the attached metal library. For convenience, we provide a 
+:meth:`mlx_build_metallib` function that builds a ``.metallib`` target given 
+sources, headers, destinations, etc. (defined in ``cmake/extension.cmake`` and 
+automatically imported with MLX package). 
 
-Here is what that looks like in practice:
+Here is what that looks like in practice!
 
 .. code-block:: cmake
 
@@ -736,29 +779,27 @@ Here is what that looks like in practice:
 
     endif()
 
-Finally, we build the nanobind_ bindings
+Finally, we build the Pybind11_ bindings
 
 .. code-block:: cmake
 
-    nanobind_add_module(
-      _ext
-      NB_STATIC STABLE_ABI LTO NOMINSIZE
-      NB_DOMAIN mlx
-      ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
+    pybind11_add_module(
+        mlx_sample_extensions
+        ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
     )
-    target_link_libraries(_ext PRIVATE mlx_ext)
+    target_link_libraries(mlx_sample_extensions PRIVATE mlx_ext)
 
     if(BUILD_SHARED_LIBS)
-      target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
+        target_link_options(mlx_sample_extensions PRIVATE -Wl,-rpath,@loader_path)
     endif()
 
 Building with ``setuptools``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 Once we have set out the CMake build rules as described above, we can use the
-build utilities defined in :mod:`mlx.extension`:
+build utilities defined in :mod:`mlx.extension` for a simple build process. 
 
-.. code-block:: python
+.. code-block:: python 
 
     from mlx import extension
     from setuptools import setup
@@ -768,50 +809,48 @@ build utilities defined in :mod:`mlx.extension`:
             name="mlx_sample_extensions",
             version="0.0.0",
             description="Sample C++ and Metal extensions for MLX primitives.",
-            ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")],
+            ext_modules=[extension.CMakeExtension("mlx_sample_extensions")],
             cmdclass={"build_ext": extension.CMakeBuild},
-            packages=["mlx_sample_extensions"],
-            package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
-            extras_require={"dev":[]},
+            packages = ["mlx_sample_extensions"],
+            package_dir = {"": "mlx_sample_extensions"},
+            package_data = {"mlx_sample_extensions" : ["*.so", "*.dylib", "*.metallib"]},
             zip_safe=False,
-            python_requires=">=3.8",
+            python_requires=">=3.7",
         )
 
 .. note::
     We treat ``extensions/mlx_sample_extensions`` as the package directory
     even though it only contains a ``__init__.py`` to ensure the following:
+    
+    * :mod:`mlx.core` is always imported before importing  :mod:`mlx_sample_extensions`
+    * The C++ extension library and the metal library are co-located with the python 
+      bindings and copied together if the package is installed 
 
-    * :mod:`mlx.core` must be imported before importing :mod:`_ext`
-    * The C++ extension library and the metal library are co-located with the python
-      bindings and copied together if the package is installed
-
-To build the package, first install the build dependencies with ``pip install
--r requirements.txt``.  You can then build inplace for development using
+You can build inplace for development using
 ``python setup.py build_ext -j8 --inplace`` (in ``extensions/``)
 
-This results in the directory structure:
+This will result in a directory structure as follows:
 
 | extensions
 | ├── mlx_sample_extensions
 | │   ├── __init__.py
 | │   ├── libmlx_ext.dylib # C++ extension library
 | │   ├── mlx_ext.metallib # Metal library
-| │   └── _ext.cpython-3x-darwin.so # Python Binding
+| │   └── mlx_sample_extensions.cpython-3x-darwin.so # Python Binding
 | ...
 
-When you try to install using the command ``python -m pip install .`` (in
-``extensions/``), the package will be installed with the same structure as
-``extensions/mlx_sample_extensions`` and the C++ and Metal library will be
-copied along with the Python binding since they are specified as
-``package_data``.
+When you try to install using the command ``python -m pip install .`` 
+(in ``extensions/``), the package will be installed with the same structure as 
+``extensions/mlx_sample_extensions`` and the C++ and metal library will be 
+copied along with the python binding since they are specified as ``package_data``.
 
 Usage
 -----
 
-After installing the extension as described above, you should be able to simply
-import the Python package and play with it as you would any other MLX operation.
+After installing the extension as described above, you should be able to simply 
+import the python package and play with it as you would any other MLX operation!
 
-Let's look at a simple script and its results:
+Let's looks at a simple script and it's results!
 
 .. code-block:: python
 
@@ -824,7 +863,7 @@ Let's look at a simple script and its results:
 
     print(f"c shape: {c.shape}")
     print(f"c dtype: {c.dtype}")
-    print(f"c correct: {mx.all(c == 6.0).item()}")
+    print(f"c correctness: {mx.all(c == 6.0).item()}")
 
 Output:
 
@@ -835,12 +874,12 @@ Output:
     c correctness: True
 
 Results
-^^^^^^^
+^^^^^^^^^^^^^^^^
 
-Let's run a quick benchmark and see how our new ``axpby`` operation compares
-with the naive :meth:`simple_axpby` we first defined on the CPU.
+Let's run a quick benchmark and see how our new ``axpby`` operation compares 
+with the naive :meth:`simple_axpby` we defined at first on the CPU. 
 
-.. code-block:: python
+.. code-block:: python 
 
     import mlx.core as mx
     from mlx_sample_extensions import axpby
@@ -859,7 +898,7 @@ with the naive :meth:`simple_axpby` we first defined on the CPU.
     alpha = 4.0
     beta = 2.0
 
-    mx.eval(x, y)
+    mx.eval((x, y))
 
     def bench(f):
         # Warm up
@@ -880,23 +919,30 @@ with the naive :meth:`simple_axpby` we first defined on the CPU.
 
     print(f"Simple axpby: {simple_time:.3f} s | Custom axpby: {custom_time:.3f} s")
 
-The results are ``Simple axpby: 0.114 s | Custom axpby: 0.109 s``. We see
-modest improvements right away!
+Results:
+
+.. code-block::
+
+    Simple axpby: 0.114 s | Custom axpby: 0.109 s
+
+We see some modest improvements right away! 
 
 This operation is now good to be used to build other operations, in
 :class:`mlx.nn.Module` calls, and also as a part of graph transformations like
-:meth:`grad`.
+:meth:`grad`!
 
 Scripts
 -------
 
 .. admonition:: Download the code
 
-   The full example code is available in `mlx <https://github.com/ml-explore/mlx/tree/main/examples/extensions/>`_.
+   The full example code is available in `mlx <code>`_.
+
+.. code: `https://github.com/ml-explore/mlx/tree/main/examples/extensions/`_
 
 .. _Accelerate: https://developer.apple.com/documentation/accelerate/blas?language=objc
 .. _Metal: https://developer.apple.com/documentation/metal?language=objc
 .. _Metal-cpp: https://developer.apple.com/metal/cpp/
 .. _`Metal Specification`: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
 .. _`Metal Example`: https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu?language=objc
-.. _nanobind: https://nanobind.readthedocs.io/en/latest/
+.. _PyBind11: https://pybind11.readthedocs.io/en/stable/

docs/src/dev/metal_debugger.rst

@@ -1,68 +0,0 @@
-Metal Debugger
-==============
-
-.. currentmodule:: mlx.core
-
-Profiling is a key step for performance optimization. You can build MLX with
-the ``MLX_METAL_DEBUG`` option to improve the Metal debugging and
-optimization workflow. The ``MLX_METAL_DEBUG`` debug option:
-
-* Records source during Metal compilation, for later inspection while
-  debugging.
-* Labels Metal objects such as command queues, improving capture readability.
-
-To build with debugging enabled in Python prepend
-``CMAKE_ARGS="-DMLX_METAL_DEBUG=ON"`` to the build call.
-
-The :func:`metal.start_capture` function initiates a capture of all MLX GPU
-work.
-
-.. note::
-
-   To capture a GPU trace you must run the application with
-   ``MTL_CAPTURE_ENABLED=1``.
-
-.. code-block:: python
-
-    import mlx.core as mx
-
-    a = mx.random.uniform(shape=(512, 512))
-    b = mx.random.uniform(shape=(512, 512))
-    mx.eval(a, b)
-
-    trace_file = "mlx_trace.gputrace"
-
-    # Make sure to run with MTL_CAPTURE_ENABLED=1 and
-    # that the path trace_file does not already exist.
-    mx.metal.start_capture(trace_file)
-
-    for _ in range(10):
-      mx.eval(mx.add(a, b))
-
-    mx.metal.stop_capture()
-
-You can open and replay the GPU trace in Xcode. The ``Dependencies`` view
-has a great overview of all operations. Checkout the `Metal debugger
-documentation`_ for more information.
-
-.. image:: ../_static/metal_debugger/capture.png
-    :class: dark-light
-
-Xcode Workflow
---------------
-
-You can skip saving to a path by running within Xcode. First, generate an
-Xcode project using CMake.
-
-.. code-block::
-
-    mkdir build && cd build
-    cmake .. -DMLX_METAL_DEBUG=ON -G Xcode
-    open mlx.xcodeproj
-
-Select the ``metal_capture`` example schema and run.
-
-.. image:: ../_static/metal_debugger/schema.png
-    :class: dark-light
-
-.. _`Metal debugger documentation`: https://developer.apple.com/documentation/xcode/metal-debugger

docs/src/examples/llama-inference.rst

@@ -15,7 +15,7 @@ module to concisely define the model architecture.
 Attention layer
 ^^^^^^^^^^^^^^^^
 
-We will start with the Llama attention layer which notably uses the RoPE
+We will start with the llama attention layer which notably uses the RoPE
 positional encoding. [1]_ In addition, our attention layer will optionally use a
 key/value cache that will be concatenated with the provided keys and values to
 support efficient inference.

docs/src/examples/mlp.rst

@@ -64,7 +64,7 @@ set:
 Next, setup the problem parameters and load the data. To load the data, you need our
 `mnist data loader
 <https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py>`_, which
-we will import as ``mnist``.
+we will import as `mnist`.
 
 .. code-block:: python
 

docs/src/index.rst

@@ -43,7 +43,6 @@ are the CPU and GPU.
    usage/function_transforms
    usage/compile
    usage/numpy
-   usage/distributed
    usage/using_streams
 
 .. toctree::
@@ -59,18 +58,15 @@ are the CPU and GPU.
    :maxdepth: 1
 
    python/array
-   python/data_types
    python/devices_and_streams
    python/ops
    python/random
    python/transforms
-   python/fast
    python/fft
    python/linalg
    python/metal
    python/nn
    python/optimizers
-   python/distributed
    python/tree_utils
 
 .. toctree::
@@ -84,4 +80,3 @@ are the CPU and GPU.
    :maxdepth: 1
 
    dev/extensions
-   dev/metal_debugger

docs/src/install.rst

@@ -15,10 +15,10 @@ To install from PyPI you must meet the following requirements:
 
 - Using an M series chip (Apple silicon)
 - Using a native Python >= 3.8
-- macOS >= 13.5
+- macOS >= 13.3
 
 .. note::
-    MLX is only available on devices running macOS >= 13.5
+    MLX is only available on devices running macOS >= 13.3 
     It is highly recommended to use macOS 14 (Sonoma)
 
 
@@ -54,7 +54,7 @@ Build Requirements
 
 - A C++ compiler with C++17 support (e.g. Clang >= 5.0)
 - `cmake <https://cmake.org/>`_ -- version 3.24 or later, and ``make``
-- Xcode >= 15.0 and macOS SDK >= 14.0
+- Xcode >= 14.3 (Xcode >= 15.0 for macOS 14 and above)
 
 .. note::
    Ensure your shell environment is native ``arm``, not ``x86`` via Rosetta. If
@@ -70,36 +70,39 @@ 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
 
-Then simply build and install MLX using pip:
+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:
 
 .. code-block:: shell
 
-  CMAKE_BUILD_PARALLEL_LEVEL="" pip install .
+    pip install "pybind11[global]"
+    conda install pybind11
+    brew install pybind11
 
-For developing, install the package with development dependencies, and use an
-editable install:
+Then simply build and install it using pip:
 
 .. code-block:: shell
 
-  CMAKE_BUILD_PARALLEL_LEVEL="" pip install -e ".[dev]"
+   env CMAKE_BUILD_PARALLEL_LEVEL="" pip install .
 
-Once the development dependencies are installed, you can build faster with:
+For developing use an editable install:
 
 .. code-block:: shell
 
- CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py build_ext -j --inplace
+  env CMAKE_BUILD_PARALLEL_LEVEL="" pip install -e .
 
-Run the tests with:
+To make sure the install is working run the tests with:
 
 .. code-block:: shell
 
+  pip install ".[testing]"
   python -m unittest discover python/tests
 
-Optional: Install stubs to enable auto completions and type checking from your
-IDE:
+Optional: Install stubs to enable auto completions and type checking from your IDE:
 
 .. code-block:: shell
 
+  pip install ".[dev]"
   python setup.py generate_stubs
 
 C++ API
@@ -120,7 +123,7 @@ Create a build directory and run CMake and make:
 .. code-block:: shell
 
    mkdir -p build && cd build
-   cmake .. && make -j
+   cmake .. && make -j 
 
 Run tests with:
 
@@ -139,7 +142,7 @@ directory as the executable statically linked to ``libmlx.a`` or the
 preprocessor constant ``METAL_PATH`` should be defined at build time and it
 should point to the path to the built metal library.
 
-.. list-table:: Build Options
+.. list-table:: Build Options 
    :widths: 25 8
    :header-rows: 1
 
@@ -153,67 +156,31 @@ should point to the path to the built metal library.
      - OFF
    * - MLX_BUILD_METAL
      - ON
-   * - MLX_BUILD_CPU
-     - ON
    * - MLX_BUILD_PYTHON_BINDINGS
      - OFF
-   * - MLX_METAL_DEBUG
-     - OFF
-   * - MLX_BUILD_SAFETENSORS
-     - ON
-   * - MLX_BUILD_GGUF
-     - ON
-   * - MLX_METAL_JIT
-     - OFF
+
 
 .. note::
 
-    If you have multiple Xcode installations and wish to use
-    a specific one while building, you can do so by adding the
-    following environment variable before building
+    If you have multiple Xcode installations and wish to use 
+    a specific one while building, you can do so by adding the 
+    following environment variable before building 
 
     .. code-block:: shell
 
       export DEVELOPER_DIR="/path/to/Xcode.app/Contents/Developer/"
 
-    Further, you can use the following command to find out which
+    Further, you can use the following command to find out which 
     macOS SDK will be used
 
     .. code-block:: shell
 
       xcrun -sdk macosx --show-sdk-version
 
-Binary Size Minimization
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-To produce a smaller binary use the CMake flags ``CMAKE_BUILD_TYPE=MinSizeRel``
-and ``BUILD_SHARED_LIBS=ON``.
-
-The MLX CMake build has several additional options to make smaller binaries.
-For example, if you don't need the CPU backend or support for safetensors and
-GGUF, you can do:
-
-.. code-block:: shell
-
-  cmake .. \
-    -DCMAKE_BUILD_TYPE=MinSizeRel \
-    -DBUILD_SHARED_LIBS=ON \
-    -DMLX_BUILD_CPU=OFF \
-    -DMLX_BUILD_SAFETENSORS=OFF \
-    -DMLX_BUILD_GGUF=OFF \
-    -DMLX_METAL_JIT=ON
-
-THE ``MLX_METAL_JIT`` flag minimizes the size of the MLX Metal library which
-contains pre-built GPU kernels. This substantially reduces the size of the
-Metal library by run-time compiling kernels the first time they are used in MLX
-on a given machine. Note run-time compilation incurs a cold-start cost which can
-be anwywhere from a few hundred millisecond to a few seconds depending on the
-application. Once a kernel is compiled, it will be cached by the system. The
-Metal kernel cache persists accross reboots.
-
 Troubleshooting
 ^^^^^^^^^^^^^^^
 
+
 Metal not found
 ~~~~~~~~~~~~~~~
 
@@ -235,7 +202,7 @@ Then set the active developer directory:
 
   sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
 
-x86 Shell
+x86 Shell 
 ~~~~~~~~~
 
 .. _build shell:

docs/src/python/array.rst

@@ -10,39 +10,27 @@ Array
 
     array
     array.astype
-    array.at
     array.item
     array.tolist
     array.dtype
-    array.itemsize
-    array.nbytes
     array.ndim
     array.shape
     array.size
+    Dtype
     array.abs
     array.all
     array.any
     array.argmax
     array.argmin
-    array.conj
     array.cos
-    array.cummax
-    array.cummin
-    array.cumprod
-    array.cumsum
-    array.diag
-    array.diagonal
+    array.dtype
     array.exp
-    array.flatten
     array.log
-    array.log10
     array.log1p
-    array.log2
     array.logsumexp
     array.max
     array.mean
     array.min
-    array.moveaxis
     array.prod
     array.reciprocal
     array.reshape
@@ -52,10 +40,7 @@ Array
     array.split
     array.sqrt
     array.square
-    array.squeeze
-    array.swapaxes
     array.sum
     array.transpose
     array.T
     array.var
-    array.view

docs/src/python/data_types.rst

@@ -1,5 +1,7 @@
 .. _data_types:
 
+:orphan:
+
 Data Types
 ==========
 
@@ -42,27 +44,9 @@ The default floating point type is ``float32`` and the default integer type is
    * - ``int64``
      - 8 
      - 64-bit signed integer 
-   * - ``bfloat16``
-     - 2 
-     - 16-bit brain float (e8, m7)
    * - ``float16``
      - 2 
-     - 16-bit IEEE float (e5, m10)
+     - 16-bit float, only available with `ARM C language extensions <https://developer.arm.com/documentation/101028/0012/3--C-language-extensions?lang=en>`_
    * - ``float32``
      - 4 
      - 32-bit float
-   * - ``complex64``
-     - 8 
-     - 64-bit complex float
-
-
-Data type are aranged in a hierarchy. See the :obj:`DtypeCategory` object
-documentation for more information. Use :func:`issubdtype` to determine if one
-``dtype`` (or category) is a subtype of another category.
-
-.. autosummary::
-   :toctree: _autosummary
-
-   Dtype
-   DtypeCategory
-   issubdtype

docs/src/python/devices_and_streams.rst

@@ -16,4 +16,3 @@ Devices and Streams
    new_stream
    set_default_stream
    stream
-   synchronize

docs/src/python/distributed.rst

@@ -1,19 +0,0 @@
-.. _distributed:
-
-.. currentmodule:: mlx.core.distributed
-
-Distributed Communication
-==========================
-
-MLX provides a distributed communication package using MPI. The MPI library is
-loaded at runtime; if MPI is available then distributed communication is also
-made available.
-
-.. autosummary::
-   :toctree: _autosummary
-
-    Group
-    is_available
-    init
-    all_sum
-    all_gather

docs/src/python/fast.rst

@@ -1,14 +0,0 @@
-.. _fast:
-
-Fast
-====
-
-.. currentmodule:: mlx.core.fast
-
-.. autosummary:: 
-  :toctree: _autosummary
-
-  rms_norm
-  layer_norm
-  rope
-  scaled_dot_product_attention

docs/src/python/linalg.rst

@@ -8,10 +8,5 @@ Linear Algebra
 .. autosummary:: 
    :toctree: _autosummary 
 
-    inv
-    tri_inv
     norm
-    cholesky
-    cholesky_inv
     qr
-    svd

docs/src/python/metal.rst

@@ -3,17 +3,12 @@ Metal
 
 .. currentmodule:: mlx.core.metal
 
-.. autosummary::
+.. autosummary:: 
   :toctree: _autosummary
 
   is_available
-  device_info
   get_active_memory
   get_peak_memory
-  reset_peak_memory
   get_cache_memory
   set_memory_limit
   set_cache_limit
-  clear_cache
-  start_capture
-  stop_capture

docs/src/python/nn.rst

@@ -173,7 +173,7 @@ In detail:
    :toctree: _autosummary
 
    value_and_grad
-   quantize
+   checkpoint
 
 .. toctree::
 

docs/src/python/nn/functions.rst

@@ -17,8 +17,6 @@ simple functions.
    gelu_approx
    gelu_fast_approx
    glu
-   hard_shrink
-   hard_tanh
    hardswish
    leaky_relu
    log_sigmoid
@@ -31,7 +29,6 @@ simple functions.
    sigmoid
    silu
    softmax
-   softmin
    softplus
    softshrink
    step

docs/src/python/nn/layers.rst

@@ -15,45 +15,29 @@ Layers
    BatchNorm
    Conv1d
    Conv2d
-   Conv3d
    Dropout
    Dropout2d
    Dropout3d
    Embedding
    GELU
-   GLU
    GroupNorm
-   GRU
-   HardShrink
-   HardTanh
-   Hardswish
    InstanceNorm
    LayerNorm
-   LeakyReLU
    Linear
-   LSTM
    MaxPool1d
    MaxPool2d
    Mish
    MultiHeadAttention
    PReLU
-   QuantizedEmbedding
    QuantizedLinear
    RMSNorm
    ReLU
-   ReLU6
-   RNN
    RoPE
    SELU
    Sequential
    SiLU
    SinusoidalPositionalEncoding
-   Softmin
    Softshrink
-   Softsign
-   Softmax
-   Softplus
    Step
-   Tanh
    Transformer
-   Upsample
+   Upsample

docs/src/python/nn/module.rst

@@ -30,7 +30,6 @@ Module
       Module.named_modules
       Module.parameters
       Module.save_weights
-      Module.set_dtype
       Module.train
       Module.trainable_parameters
       Module.unfreeze

docs/src/python/ops.rst

@@ -5,14 +5,13 @@ Operations
 
 .. currentmodule:: mlx.core
 
-.. autosummary::
+.. autosummary:: 
   :toctree: _autosummary
 
    abs
    add
-   addmm
    all
-   allclose
+   allclose 
    any
    arange
    arccos
@@ -20,70 +19,49 @@ Operations
    arcsin
    arcsinh
    arctan
-   arctan2
    arctanh
    argmax
    argmin
    argpartition
    argsort
    array_equal
-   as_strided
    atleast_1d
    atleast_2d
    atleast_3d
-   bitwise_and
-   bitwise_or
-   bitwise_xor
-   block_masked_mm
    broadcast_to
    ceil
    clip
    concatenate
-   conj
-   conjugate
    convolve
    conv1d
    conv2d
    conv_general
    cos
    cosh
-   cummax
-   cummin
-   cumprod
-   cumsum
-   degrees
    dequantize
    diag
    diagonal
    divide
    divmod
-   einsum
-   einsum_path
    equal
    erf
    erfinv
    exp
-   expm1
    expand_dims
    eye
    flatten
    floor
    floor_divide
    full
-   gather_mm
-   gather_qmm
    greater
    greater_equal
-   hadamard_transform
    identity
    inner
    isclose
-   isinf
    isnan
-   isneginf
    isposinf
-   issubdtype
-   left_shift
+   isneginf
+   isinf
    less
    less_equal
    linspace
@@ -101,29 +79,22 @@ Operations
    max
    maximum
    mean
-   meshgrid
    min
    minimum
    moveaxis
    multiply
-   nan_to_num
    negative
-   not_equal
    ones
    ones_like
    outer
    partition
    pad
-   power
    prod
    quantize
    quantized_matmul
-   radians
    reciprocal
-   remainder
    repeat
    reshape
-   right_shift
    round
    rsqrt
    save
@@ -142,7 +113,6 @@ Operations
    square
    squeeze
    stack
-   std
    stop_gradient
    subtract
    sum
@@ -154,13 +124,11 @@ Operations
    tensordot
    tile
    topk
-   trace
    transpose
    tri
    tril
    triu
    var
-   view
    where
    zeros
    zeros_like

docs/src/python/optimizers.rst

@@ -1,7 +1,5 @@
 .. _optimizers:
 
-.. currentmodule:: mlx.optimizers
-
 Optimizers
 ==========
 
@@ -31,48 +29,8 @@ model's parameters and the **optimizer state**.
             # Compute the new parameters but also the optimizer state.
             mx.eval(model.parameters(), optimizer.state)
 
-Saving and Loading
-------------------
-
-To serialize an optimizer, save its state. To load an optimizer, load and set
-the saved state. Here's a simple example:
-
-.. code-block:: python
-
-   import mlx.core as mx
-   from mlx.utils import tree_flatten, tree_unflatten
-   import mlx.optimizers as optim
-
-   optimizer = optim.Adam(learning_rate=1e-2)
-
-   # Perform some updates with the optimizer
-   model = {"w" : mx.zeros((5, 5))}
-   grads = {"w" : mx.ones((5, 5))}
-   optimizer.update(model, grads)
-
-   # Save the state
-   state = tree_flatten(optimizer.state)
-   mx.save_safetensors("optimizer.safetensors", dict(state))
-
-   # Later on, for example when loading from a checkpoint,
-   # recreate the optimizer and load the state
-   optimizer = optim.Adam(learning_rate=1e-2)
-
-   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
-   optimizer.state = state
-
-Note, not every optimizer configuation parameter is saved in the state. For
-example, for Adam the learning rate is saved but the ``betas`` and ``eps``
-parameters are not. A good rule of thumb is if the parameter can be scheduled
-then it will be included in the optimizer state.
-
 .. toctree::
 
    optimizers/optimizer
    optimizers/common_optimizers
    optimizers/schedulers
-
-.. autosummary::
-   :toctree: _autosummary
-
-   clip_grad_norm

docs/src/python/random.rst

@@ -38,10 +38,8 @@ we use a splittable version of Threefry, which is a counter-based PRNG.
    gumbel
    key
    normal
-   multivariate_normal
    randint
    seed
    split
    truncated_normal
    uniform
-   laplace

docs/src/python/transforms.rst

@@ -10,7 +10,6 @@ Transforms
 
    eval
    compile
-   custom_function
    disable_compile
    enable_compile
    grad
@@ -18,3 +17,4 @@ Transforms
    jvp
    vjp
    vmap
+   checkpoint

docs/src/python/tree_utils.rst

@@ -19,5 +19,3 @@ return python trees will be using the default python ``dict``, ``list`` and
    tree_flatten
    tree_unflatten
    tree_map
-   tree_map_with_path
-   tree_reduce

docs/src/usage/distributed.rst

@@ -1,166 +0,0 @@
-.. _usage_distributed:
-
-Distributed Communication
-=========================
-
-.. currentmodule:: mlx.core.distributed
-
-MLX utilizes `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ to
-provide distributed communication operations that allow the computational cost
-of training or inference to be shared across many physical machines. You can
-see a list of the supported operations in the :ref:`API docs<distributed>`.
-
-.. note::
-   A lot of operations may not be supported or not as fast as they should be.
-   We are adding more and tuning the ones we have as we are figuring out the
-   best way to do distributed computing on Macs using MLX.
-
-Getting Started
----------------
-
-MLX already comes with the ability to "talk" to MPI if it is installed on the
-machine. The minimal distributed program in MLX is as simple as:
-
-.. code:: python
-
-    import mlx.core as mx
-
-    world = mx.distributed.init()
-    x = mx.distributed.all_sum(mx.ones(10))
-    print(world.rank(), x)
-
-The program above sums the array ``mx.ones(10)`` across all
-distributed processes. If simply run with ``python``, however, only one
-process is launched and no distributed communication takes place.
-
-To launch the program in distributed mode we need to use ``mpirun`` or
-``mpiexec`` depending on the MPI installation. The simplest possible way is the
-following:
-
-.. code:: shell
-
-    $ mpirun -np 2 python test.py
-    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-
-The above launches two processes on the same (local) machine and we can see
-both standard output streams. The processes send the array of 1s to each other
-and compute the sum which is printed. Launching with ``mpirun -np 4 ...`` would
-print 4 etc.
-
-Installing MPI
----------------
-
-MPI can be installed with Homebrew, using the Anaconda package manager or
-compiled from source. Most of our testing is done using ``openmpi`` installed
-with the Anaconda package manager as follows:
-
-.. code:: shell
-
-    $ conda install openmpi
-
-Installing with Homebrew may require specifying the location of ``libmpi.dyld``
-so that MLX can find it and load it at runtime. This can simply be achieved by
-passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun``.
-
-.. code:: shell
-
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
-
-Setting up Remote Hosts
------------------------
-
-MPI can automatically connect to remote hosts and set up the communication over
-the network if the remote hosts can be accessed via ssh. A good checklist to
-debug connectivity issues is the following:
-
-* ``ssh hostname`` works from all machines to all machines without asking for
-  password or host confirmation
-* ``mpirun`` is accessible on all machines. You can call ``mpirun`` using its
-  full path to force all machines to use a specific path.
-* Ensure that the ``hostname`` used by MPI is the one that you have configured
-  in the ``.ssh/config`` files on all machines.
-
-.. note::
-  For an example hostname ``foo.bar.com`` MPI can use only ``foo`` as
-  the hostname passed to ssh if the current hostname matches ``*.bar.com``.
-
-An easy way to pass the host names to MPI is using a host file. A host file
-looks like the following, where ``host1`` and ``host2`` should be the fully
-qualified domain names or IPs for these hosts.
-
-.. code::
-
-    host1 slots=1
-    host2 slots=1
-
-When using MLX, it is very likely that you want to use 1 slot per host, ie one
-process per host.  The hostfile also needs to contain the current
-host if you want to run on the local host. Passing the host file to
-``mpirun`` is simply done using the ``--hostfile`` command line argument.
-
-Training Example
-----------------
-
-In this section we will adapt an MLX training loop to support data parallel
-distributed training. Namely, we will average the gradients across a set of
-hosts before applying them to the model.
-
-Our training loop looks like the following code snippet if we omit the model,
-dataset and optimizer initialization.
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())
-
-All we have to do to average the gradients across machines is perform an
-:func:`all_sum` and divide by the size of the :class:`Group`. Namely we
-have to :func:`mlx.utils.tree_map` the gradients with following function.
-
-.. code:: python
-
-    def all_avg(x):
-        return mx.distributed.all_sum(x) / mx.distributed.init().size()
-
-Putting everything together our training loop step looks as follows with
-everything else remaining the same.
-
-.. code:: python
-
-    from mlx.utils import tree_map
-
-    def all_reduce_grads(grads):
-        N = mx.distributed.init()
-        if N == 1:
-            return grads
-        return tree_map(
-                lambda x: mx.distributed.all_sum(x) / N,
-                grads)
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = all_reduce_grads(grads)  # <--- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-Tuning All Reduce
------------------
-
-We are working on improving the performance of all reduce on MLX but for now
-the two main things one can do to extract the most out of distributed training with MLX are:
-
-1. Perform a few large reductions instead of many small ones to improve
-   bandwidth and latency
-2. Pass ``--mca btl_tcp_links 4`` to ``mpirun`` to configure it to use 4 tcp
-   connections between each host to improve bandwidth

docs/src/usage/function_transforms.rst

@@ -40,7 +40,7 @@ getting higher order derivatives.
 
 Any of the MLX function transformations can be composed in any order to any
 depth. See the following sections for more information on :ref:`automatic
-differentiation <auto diff>` and :ref:`automatic vectorization <vmap>`.
+differentiaion <auto diff>` and :ref:`automatic vectorization <vmap>`.
 For more information on :func:`compile` see the :ref:`compile documentation <compile>`.
 
 

docs/src/usage/lazy_evaluation.rst

@@ -18,7 +18,7 @@ describe below.
 Transforming Compute Graphs
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Lazy evaluation lets us record a compute graph without actually doing any
+Lazy evaluation let's us record a compute graph without actually doing any
 computations. This is useful for function transformations like :func:`grad` and
 :func:`vmap` and graph optimizations.
 

docs/src/usage/numpy.rst

@@ -3,11 +3,7 @@
 Conversion to NumPy and Other Frameworks
 ========================================
 
-MLX array supports conversion between other frameworks with either:  
-
-* The `Python Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_. 
-* `DLPack <https://dmlc.github.io/dlpack/latest/>`_.  
-
+MLX array implements the `Python Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_.
 Let's convert an array to NumPy and back.
 
 .. code-block:: python

docs/src/usage/saving_and_loading.rst

@@ -49,7 +49,7 @@ it will be added. You can load the array with:
 
 .. code-block:: shell
 
-   >>> mx.load("array.npy")
+   >>> mx.load("array.npy", a)
    array([1], dtype=float32)
 
 Here's an example of saving several arrays to a single file:

examples/cpp/CMakeLists.txt

@@ -8,5 +8,3 @@ endfunction(build_example)
 build_example(tutorial.cpp)
 build_example(linear_regression.cpp)
 build_example(logistic_regression.cpp)
-build_example(metal_capture.cpp)
-build_example(distributed.cpp)

examples/cpp/distributed.cpp

@@ -1,22 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include <iostream>
-
-#include "mlx/mlx.h"
-
-using namespace mlx::core;
-
-int main() {
-  if (!distributed::is_available()) {
-    std::cout << "No communication backend found" << std::endl;
-    return 1;
-  }
-
-  auto global_group = distributed::init();
-  std::cout << global_group.rank() << " / " << global_group.size() << std::endl;
-
-  array x = ones({10});
-  array out = distributed::all_sum(x, global_group);
-
-  std::cout << out << std::endl;
-}

examples/cpp/metal_capture.cpp

@@ -1,31 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include <cassert>
-#include <iostream>
-
-#include "mlx/mlx.h"
-
-using namespace mlx::core;
-
-int main() {
-  // To use Metal debugging and profiling:
-  // 1. Build with the MLX_METAL_DEBUG CMake option (i.e. -DMLX_METAL_DEBUG=ON).
-  // 2. Run with MTL_CAPTURE_ENABLED=1.
-  metal::start_capture("mlx_trace.gputrace");
-
-  // Start at index two because the default GPU and CPU streams have indices
-  // zero and one, respectively. This naming matches the label assigned to each
-  // stream's command queue.
-  auto s2 = new_stream(Device::gpu);
-  auto s3 = new_stream(Device::gpu);
-
-  auto a = arange(1.f, 10.f, 1.f, float32, s2);
-  auto b = arange(1.f, 10.f, 1.f, float32, s3);
-  auto x = add(a, a, s2);
-  auto y = add(b, b, s3);
-
-  // The multiply will happen on the default stream.
-  std::cout << multiply(x, y) << std::endl;
-
-  metal::stop_capture();
-}

examples/cpp/tutorial.cpp

@@ -89,8 +89,8 @@ void automatic_differentiation() {
   // dfdx is 2 * x
 
   // Get the second derivative by composing grad with grad
-  auto d2fdx2 = grad(grad(fn))(x);
-  // d2fdx2 is 2
+  auto df2dx2 = grad(grad(fn))(x);
+  // df2dx2 is 2
 }
 
 int main() {

examples/extensions/CMakeLists.txt

@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.27)
 
-project(_ext LANGUAGES CXX)
+project(mlx_sample_extensions LANGUAGES CXX)
 
 # ----------------------------- Setup -----------------------------
 set(CMAKE_CXX_STANDARD 17)
@@ -11,12 +11,8 @@ option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
 
 # ----------------------------- Dependencies -----------------------------
 find_package(MLX 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)
+find_package(Python COMPONENTS Interpreter Development)
+find_package(pybind11 CONFIG REQUIRED)
 
 # ----------------------------- Extensions -----------------------------
 
@@ -42,6 +38,7 @@ target_link_libraries(mlx_ext PUBLIC mlx)
 
 # Build metallib
 if(MLX_BUILD_METAL)
+
   mlx_build_metallib(
     TARGET mlx_ext_metallib
     TITLE mlx_ext
@@ -57,15 +54,13 @@ if(MLX_BUILD_METAL)
 
 endif()
 
-# ----------------------------- Python Bindings -----------------------------
-nanobind_add_module(
-  _ext
-  NB_STATIC STABLE_ABI LTO NOMINSIZE
-  NB_DOMAIN mlx 
+# ----------------------------- Pybind -----------------------------
+pybind11_add_module(
+  mlx_sample_extensions
   ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
 )
-target_link_libraries(_ext PRIVATE mlx_ext)
+target_link_libraries(mlx_sample_extensions PRIVATE mlx_ext)
 
 if(BUILD_SHARED_LIBS)
-  target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
+  target_link_options(mlx_sample_extensions PRIVATE -Wl,-rpath,@loader_path)
 endif()

examples/extensions/README.md

@@ -1,24 +0,0 @@
-
-## Build
-
-```
-pip install -e .
-```
-
-For faster builds during development, you can also pre-install the requirements:
-
-```
-pip install -r requirements.txt
-```
-
-And then run:
-
-```
-python setup.py build_ext -j8 --inplace
-```
-
-## Test
-
-```
-python test.py
-```

examples/extensions/axpby/axpby.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <cassert>
 #include <iostream>
@@ -43,7 +43,7 @@ array axpby(
   auto promoted_dtype = promote_types(x.dtype(), y.dtype());
 
   // Upcast to float32 for non-floating point inputs x and y
-  auto out_dtype = issubdtype(promoted_dtype, float32)
+  auto out_dtype = is_floating_point(promoted_dtype)
       ? promoted_dtype
       : promote_types(promoted_dtype, float32);
 
@@ -61,7 +61,7 @@ array axpby(
       /* const std::vector<int>& shape = */ out_shape,
       /* Dtype dtype = */ out_dtype,
       /* std::unique_ptr<Primitive> primitive = */
-      std::make_shared<Axpby>(to_stream(s), alpha, beta),
+      std::make_unique<Axpby>(to_stream(s), alpha, beta),
       /* const std::vector<array>& inputs = */ broadcasted_inputs);
 }
 
@@ -106,12 +106,12 @@ void axpby_impl(
 /** Fall back implementation for evaluation on CPU */
 void Axpby::eval(
     const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    std::vector<array>& out_arr) {
+  auto out = out_arr[0];
   // Check the inputs (registered in the op while constructing the out array)
   assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
-  auto& out = outputs[0];
 
   // Dispatch to the correct dtype
   if (out.dtype() == float32) {
@@ -150,7 +150,11 @@ void axpby_impl_accelerate(
   // The data in the output array is allocated to match the strides in y
   // such that x, y, and out are contiguous in the same mode and
   // no transposition is needed
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  out.set_data(
+      allocator::malloc_or_wait(y.data_size() * out.itemsize()),
+      y.data_size(),
+      y.strides(),
+      y.flags());
 
   // We then copy over the elements using the contiguous vector specialization
   copy_inplace(y, out, CopyType::Vector);
@@ -176,11 +180,11 @@ void axpby_impl_accelerate(
 /** Evaluate primitive on CPU using accelerate specializations */
 void Axpby::eval_cpu(
     const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    std::vector<array>& outarr) {
+  auto out = outarr[0];
   assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
-  auto& out = outputs[0];
 
   // Accelerate specialization for contiguous single precision float arrays
   if (out.dtype() == float32 &&
@@ -191,7 +195,7 @@ void Axpby::eval_cpu(
   }
 
   // Fall back to common backend if specializations are not available
-  eval(inputs, outputs);
+  eval(inputs, outarr);
 }
 
 #else // Accelerate not available
@@ -199,8 +203,8 @@ void Axpby::eval_cpu(
 /** Evaluate primitive on CPU falling back to common backend */
 void Axpby::eval_cpu(
     const std::vector<array>& inputs,
-    const std::vector<array>& outputs) {
-  eval(inputs, outputs);
+    std::vector<array>& out) {
+  eval(inputs, out);
 }
 
 #endif
@@ -214,12 +218,12 @@ void Axpby::eval_cpu(
 /** Evaluate primitive on GPU */
 void Axpby::eval_gpu(
     const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    std::vector<array>& outarr) {
   // Prepare inputs
+  auto out = outarr[0];
   assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
-  auto& out = outputs[0];
 
   // Each primitive carries the stream it should execute on
   // and each stream carries its device identifiers
@@ -249,14 +253,15 @@ void Axpby::eval_gpu(
   kname << (contiguous_kernel ? "contiguous_" : "general_");
   kname << type_to_name(out);
 
-  // Make sure the metal library is available
-  d.register_library("mlx_ext");
+  // Make sure the metal library is available and look for it
+  // in the same folder as this executable if needed
+  d.register_library("mlx_ext", metal::get_colocated_mtllib_path);
 
   // Make a kernel from this metal library
   auto kernel = d.get_kernel(kname.str(), "mlx_ext");
 
   // Prepare to encode kernel
-  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
 
   // Kernel parameters are registered with buffer indices corresponding to
@@ -265,11 +270,11 @@ void Axpby::eval_gpu(
   size_t nelem = out.size();
 
   // Encode input arrays to kernel
-  compute_encoder.set_input_array(x, 0);
-  compute_encoder.set_input_array(y, 1);
+  set_array_buffer(compute_encoder, x, 0);
+  set_array_buffer(compute_encoder, y, 1);
 
   // Encode output arrays to kernel
-  compute_encoder.set_output_array(out, 2);
+  set_array_buffer(compute_encoder, out, 2);
 
   // Encode alpha and beta
   compute_encoder->setBytes(&alpha_, sizeof(float), 3);
@@ -295,7 +300,7 @@ void Axpby::eval_gpu(
 
   // Launch the grid with the given number of threads divided among
   // the given threadgroups
-  compute_encoder.dispatchThreads(grid_dims, group_dims);
+  compute_encoder->dispatchThreads(grid_dims, group_dims);
 }
 
 #else // Metal is not available
@@ -367,4 +372,4 @@ bool Axpby::is_equivalent(const Primitive& other) const {
   return alpha_ == r_other.alpha_ && beta_ == r_other.beta_;
 }
 
-} // namespace mlx::core
+} // namespace mlx::core

examples/extensions/axpby/axpby.h

@@ -33,7 +33,7 @@ array axpby(
 class Axpby : public Primitive {
  public:
   explicit Axpby(Stream stream, float alpha, float beta)
-      : Primitive(stream), alpha_(alpha), beta_(beta) {};
+      : Primitive(stream), alpha_(alpha), beta_(beta){};
 
   /**
    * A primitive must know how to evaluate itself on the CPU/GPU
@@ -42,9 +42,9 @@ class Axpby : public Primitive {
    * To avoid unnecessary allocations, the evaluation function
    * is responsible for allocating space for the array.
    */
-  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& out)
       override;
-  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& out)
       override;
 
   /** The Jacobian-vector product. */
@@ -83,7 +83,7 @@ class Axpby : public Primitive {
   float beta_;
 
   /** Fall back implementation for evaluation on CPU */
-  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+  void eval(const std::vector<array>& inputs, std::vector<array>& out);
 };
 
-} // namespace mlx::core
+} // namespace mlx::core

examples/extensions/axpby/axpby.metal

@@ -19,7 +19,7 @@ template <typename T>
     uint index [[thread_position_in_grid]]) {
   auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
   auto y_offset = elem_to_loc(index, shape, y_strides, ndim);
-  out[index] =
+  out[index] = 
       static_cast<T>(alpha) * x[x_offset] + static_cast<T>(beta) * y[y_offset];
 }
 
@@ -31,30 +31,30 @@ template <typename T>
     constant const float& alpha [[buffer(3)]],
     constant const float& beta [[buffer(4)]],
     uint index [[thread_position_in_grid]]) {
-  out[index] =
+  out[index] = 
       static_cast<T>(alpha) * x[index] + static_cast<T>(beta) * y[index];
 }
 
-#define instantiate_axpby(type_name, type)                               \
-  template [[host_name("axpby_general_" #type_name)]] [[kernel]] void    \
-  axpby_general<type>(                                                   \
-      device const type* x [[buffer(0)]],                                \
-      device const type* y [[buffer(1)]],                                \
-      device type* out [[buffer(2)]],                                    \
-      constant const float& alpha [[buffer(3)]],                         \
-      constant const float& beta [[buffer(4)]],                          \
-      constant const int* shape [[buffer(5)]],                           \
-      constant const size_t* x_strides [[buffer(6)]],                    \
-      constant const size_t* y_strides [[buffer(7)]],                    \
-      constant const int& ndim [[buffer(8)]],                            \
-      uint index [[thread_position_in_grid]]);                           \
-  template [[host_name("axpby_contiguous_" #type_name)]] [[kernel]] void \
-  axpby_contiguous<type>(                                                \
-      device const type* x [[buffer(0)]],                                \
-      device const type* y [[buffer(1)]],                                \
-      device type* out [[buffer(2)]],                                    \
-      constant const float& alpha [[buffer(3)]],                         \
-      constant const float& beta [[buffer(4)]],                          \
+#define instantiate_axpby(type_name, type)            \
+  template [[host_name("axpby_general_" #type_name)]] \
+  [[kernel]] void axpby_general<type>(                \
+      device const type* x [[buffer(0)]],             \
+      device const type* y [[buffer(1)]],             \
+      device type* out [[buffer(2)]],                 \
+      constant const float& alpha [[buffer(3)]],      \
+      constant const float& beta [[buffer(4)]],       \
+      constant const int* shape [[buffer(5)]],        \
+      constant const size_t* x_strides [[buffer(6)]], \
+      constant const size_t* y_strides [[buffer(7)]], \
+      constant const int& ndim [[buffer(8)]],         \
+      uint index [[thread_position_in_grid]]);        \
+  template [[host_name("axpby_contiguous_" #type_name)]] \
+  [[kernel]] void axpby_contiguous<type>(                \
+      device const type* x [[buffer(0)]],                \
+      device const type* y [[buffer(1)]],                \
+      device type* out [[buffer(2)]],                    \
+      constant const float& alpha [[buffer(3)]],         \
+      constant const float& beta [[buffer(4)]],          \
       uint index [[thread_position_in_grid]]);
 
 instantiate_axpby(float32, float);

examples/extensions/bindings.cpp

@@ -1,31 +1,31 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
-#include <nanobind/nanobind.h>
-#include <nanobind/stl/variant.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "axpby/axpby.h"
 
-namespace nb = nanobind;
-using namespace nb::literals;
-
+namespace py = pybind11;
+using namespace py::literals;
 using namespace mlx::core;
 
-NB_MODULE(_ext, m) {
-  m.doc() = "Sample extension for MLX";
+PYBIND11_MODULE(mlx_sample_extensions, m) {
+  m.doc() = "Sample C++ and metal extensions for MLX";
 
   m.def(
       "axpby",
       &axpby,
       "x"_a,
       "y"_a,
+      py::pos_only(),
       "alpha"_a,
       "beta"_a,
-      nb::kw_only(),
-      "stream"_a = nb::none(),
-      R"(
+      py::kw_only(),
+      "stream"_a = py::none(),
+      R"pbdoc(
         Scale and sum two vectors element-wise
         ``z = alpha * x + beta * y``
-
+        
         Follows numpy style broadcasting between ``x`` and ``y``
         Inputs are upcasted to floats if needed
 
@@ -37,5 +37,5 @@ NB_MODULE(_ext, m) {
 
         Returns:
             array: ``alpha * x + beta * y``
-      )");
-}
+      )pbdoc");
+}

examples/extensions/mlx_sample_extensions/__init__.py

@@ -2,4 +2,4 @@
 
 import mlx.core as mx
 
-from ._ext import axpby
+from .mlx_sample_extensions import *

examples/extensions/pyproject.toml

@@ -1,8 +1,3 @@
 [build-system]
-requires = [
-  "setuptools>=42",
-  "cmake>=3.24",
-  "mlx>=0.9.0",
-  "nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4",
-]
-build-backend = "setuptools.build_meta"
+requires = ["setuptools>=42", "pybind11>=2.10", "cmake>=3.24", "mlx @ git+https://github.com/mlx-explore/mlx@main"]
+build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

@@ -1,4 +0,0 @@
-setuptools>=42
-cmake>=3.24
-mlx>=0.16.2
-nanobind==2.0

examples/extensions/setup.py

@@ -1,4 +1,4 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023 Apple Inc.
 
 from setuptools import setup
 
@@ -9,11 +9,11 @@ if __name__ == "__main__":
         name="mlx_sample_extensions",
         version="0.0.0",
         description="Sample C++ and Metal extensions for MLX primitives.",
-        ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")],
+        ext_modules=[extension.CMakeExtension("mlx_sample_extensions")],
         cmdclass={"build_ext": extension.CMakeBuild},
         packages=["mlx_sample_extensions"],
+        package_dir={"": "."},
         package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
-        extras_require={"dev": []},
         zip_safe=False,
         python_requires=">=3.8",
     )

examples/extensions/test.py

@@ -1,10 +0,0 @@
-import mlx.core as mx
-from mlx_sample_extensions import axpby
-
-a = mx.ones((3, 4))
-b = mx.ones((3, 4))
-c = axpby(a, b, 4.0, 2.0, stream=mx.cpu)
-
-print(f"c shape: {c.shape}")
-print(f"c dtype: {c.dtype}")
-print(f"c correct: {mx.all(c == 6.0).item()}")

mlx/CMakeLists.txt

@@ -6,7 +6,6 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
@@ -20,17 +19,11 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h
 )
 
-if (MLX_BUILD_CPU)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
-else()
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_cpu)
-endif()
-
-add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
 if (MLX_BUILD_ACCELERATE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/accelerate)
-elseif(MLX_BUILD_CPU)
+else()
   target_sources(
     mlx
     PRIVATE

mlx/allocator.h

@@ -14,7 +14,7 @@ class Buffer {
   void* ptr_;
 
  public:
-  Buffer(void* ptr) : ptr_(ptr) {};
+  Buffer(void* ptr) : ptr_(ptr){};
 
   // Get the raw data pointer from the buffer
   void* raw_ptr();

mlx/array.cpp

@@ -1,4 +1,5 @@
 // Copyright © 2023-2024 Apple Inc.
+
 #include <functional>
 
 #include "mlx/array.h"
@@ -11,16 +12,22 @@ namespace mlx::core {
 
 namespace {
 
+std::pair<size_t, std::vector<size_t>> cum_prod(const std::vector<int>& shape) {
+  std::vector<size_t> strides(shape.size());
+  size_t cum_prod = 1;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    strides[i] = cum_prod;
+    cum_prod *= shape[i];
+  }
+  return {cum_prod, strides};
+}
+
 /** Return true if we are currently performing a function transformation in
  * order to keep the graph when evaluating tracer arrays. */
 bool in_tracing() {
   return detail::InTracing::in_tracing();
 }
 
-bool retain_graph() {
-  return detail::RetainGraph::retain_graph();
-}
-
 } // namespace
 
 array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
@@ -29,11 +36,22 @@ 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,
@@ -41,16 +59,15 @@ array::array(
           std::move(inputs))) {}
 
 std::vector<array> array::make_arrays(
-    std::vector<std::vector<int>> shapes,
+    const std::vector<std::vector<int>>& shapes,
     const std::vector<Dtype>& dtypes,
-    const std::shared_ptr<Primitive>& primitive,
+    std::shared_ptr<Primitive> primitive,
     const std::vector<array>& inputs) {
   std::vector<array> outputs;
-  for (size_t i = 0; i < shapes.size(); ++i) {
-    outputs.emplace_back(std::move(shapes[i]), dtypes[i], primitive, inputs);
+  for (int i = 0; i < shapes.size(); ++i) {
+    outputs.push_back(array(shapes[i], dtypes[i], primitive, inputs));
   }
-  // For each node in |outputs|, its siblings are the other nodes.
-  for (size_t i = 0; i < outputs.size(); ++i) {
+  for (int i = 0; i < outputs.size(); ++i) {
     auto siblings = outputs;
     siblings.erase(siblings.begin() + i);
     outputs[i].set_siblings(std::move(siblings), i);
@@ -75,10 +92,10 @@ array::array(std::initializer_list<int> data, Dtype dtype)
 /* Build an array from a shared buffer */
 array::array(
     allocator::Buffer data,
-    std::vector<int> shape,
+    const std::vector<int>& shape,
     Dtype dtype,
     deleter_t deleter)
-    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
   set_data(data, deleter);
 }
 
@@ -87,26 +104,22 @@ void array::detach() {
     s.array_desc_->inputs.clear();
     s.array_desc_->siblings.clear();
     s.array_desc_->position = 0;
+    s.array_desc_->depth = 0;
     s.array_desc_->primitive = nullptr;
   }
   array_desc_->inputs.clear();
   array_desc_->siblings.clear();
   array_desc_->position = 0;
+  array_desc_->depth = 0;
   array_desc_->primitive = nullptr;
 }
 
 void array::eval() {
-  // Ensure the array is ready to be read
-  if (status() == Status::scheduled) {
-    event().wait();
-    set_status(Status::available);
-  } else if (status() == Status::unscheduled) {
-    mlx::core::eval({*this});
-  }
+  mlx::core::eval({*this});
 }
 
 bool array::is_tracer() const {
-  return array_desc_->is_tracer && in_tracing() || retain_graph();
+  return array_desc_->is_tracer && in_tracing();
 }
 
 void array::set_data(allocator::Buffer buffer, deleter_t d) {
@@ -151,117 +164,51 @@ void array::copy_shared_buffer(const array& other) {
   copy_shared_buffer(other, other.strides(), other.flags(), other.data_size());
 }
 
-void array::move_shared_buffer(
-    array other,
-    const std::vector<size_t>& strides,
-    Flags flags,
-    size_t data_size,
-    size_t offset /* = 0 */) {
-  array_desc_->data = std::move(other.array_desc_->data);
-  array_desc_->strides = strides;
-  array_desc_->flags = flags;
-  array_desc_->data_size = data_size;
-  auto char_offset = sizeof(char) * itemsize() * offset;
-  array_desc_->data_ptr = static_cast<void*>(
-      static_cast<char*>(other.array_desc_->data_ptr) + char_offset);
-}
-
 void array::move_shared_buffer(array other) {
-  move_shared_buffer(other, other.strides(), other.flags(), other.data_size());
+  array_desc_->data = std::move(other.array_desc_->data);
+  array_desc_->strides = other.strides();
+  array_desc_->flags = other.flags();
+  array_desc_->data_size = other.data_size();
+  array_desc_->data_ptr = other.array_desc_->data_ptr;
 }
 
-array::~array() {
-  if (array_desc_ == nullptr) {
-    return;
-  }
-
-  // Ignore arrays that might be detached during eval
-  if (status() == array::Status::scheduled) {
-    return;
-  }
-
-  // Break circular reference for non-detached arrays with siblings
-  if (auto n = siblings().size(); n > 0) {
-    bool do_detach = true;
-    // If all siblings have siblings.size() references except
-    // the one we are currently destroying (which has siblings.size() + 1)
-    // then there are no more external references
-    do_detach &= (array_desc_.use_count() == (n + 1));
-    for (auto& s : siblings()) {
-      do_detach &= (s.array_desc_.use_count() == n);
-      if (!do_detach) {
-        break;
-      }
-    }
-    if (do_detach) {
-      for (auto& s : siblings()) {
-        for (auto& ss : s.siblings()) {
-          ss.array_desc_ = nullptr;
-        }
-        s.array_desc_->siblings.clear();
-      }
-    }
-  }
-}
-
-void array::ArrayDesc::init() {
-  strides.resize(shape.size());
-  size = 1;
-  for (int i = shape.size() - 1; i >= 0; --i) {
-    strides[i] = size;
-    size *= shape[i];
-  }
-  for (const auto& in : inputs) {
-    is_tracer |= in.is_tracer();
-  }
-}
-
-array::ArrayDesc::ArrayDesc(std::vector<int> shape, Dtype dtype)
-    : shape(std::move(shape)), dtype(dtype), status(Status::available) {
-  init();
+array::ArrayDesc::ArrayDesc(const std::vector<int>& shape, Dtype dtype)
+    : shape(shape), dtype(dtype) {
+  std::tie(size, strides) = cum_prod(shape);
 }
 
 array::ArrayDesc::ArrayDesc(
-    std::vector<int> shape,
+    const std::vector<int>& shape,
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
-    std::vector<array> inputs)
-    : shape(std::move(shape)),
+    const std::vector<array>& inputs)
+    : shape(shape),
       dtype(dtype),
-      status(Status::unscheduled),
       primitive(std::move(primitive)),
-      inputs(std::move(inputs)) {
-  init();
+      inputs(inputs) {
+  std::tie(size, strides) = cum_prod(this->shape);
+  for (auto& in : this->inputs) {
+    is_tracer |= in.is_tracer();
+    depth = std::max(in.graph_depth(), depth);
+  }
+  depth++;
 }
 
-array::ArrayDesc::~ArrayDesc() {
-  // When an array description is destroyed it will delete a bunch of arrays
-  // that may also destroy their corresponding descriptions and so on and so
-  // forth.
-  //
-  // This calls recursively the destructor and can result in stack overflow, we
-  // instead put them in a vector and destroy them one at a time resulting in a
-  // max stack depth of 2.
-  std::vector<std::shared_ptr<ArrayDesc>> for_deletion;
-
-  for (array& a : inputs) {
-    if (a.array_desc_.use_count() == 1) {
-      for_deletion.push_back(std::move(a.array_desc_));
-    }
-  }
-
-  while (!for_deletion.empty()) {
-    // top is going to be deleted at the end of the block *after* the arrays
-    // with inputs have been moved into the vector
-    auto top = std::move(for_deletion.back());
-    for_deletion.pop_back();
-
-    for (array& a : top->inputs) {
-      if (a.array_desc_.use_count() == 1) {
-        for_deletion.push_back(std::move(a.array_desc_));
-      }
-    }
+array::ArrayDesc::ArrayDesc(
+    std::vector<int>&& shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array>&& inputs)
+    : shape(std::move(shape)),
+      dtype(dtype),
+      primitive(std::move(primitive)),
+      inputs(std::move(inputs)) {
+  std::tie(size, strides) = cum_prod(this->shape);
+  for (auto& in : this->inputs) {
+    is_tracer |= in.is_tracer();
+    depth = std::max(in.graph_depth(), depth);
   }
+  depth++;
 }
 
 array::ArrayIterator::ArrayIterator(const array& arr, int idx)

mlx/array.h

@@ -1,6 +1,5 @@
 // Copyright © 2023 Apple Inc.
 #pragma once
-
 #include <algorithm>
 #include <cstdint>
 #include <functional>
@@ -9,7 +8,6 @@
 
 #include "mlx/allocator.h"
 #include "mlx/dtype.h"
-#include "mlx/event.h"
 
 namespace mlx::core {
 
@@ -33,7 +31,7 @@ class array {
   template <typename It>
   array(
       It data,
-      std::vector<int> shape,
+      const std::vector<int>& shape,
       Dtype dtype =
           TypeToDtype<typename std::iterator_traits<It>::value_type>());
 
@@ -49,13 +47,13 @@ class array {
   template <typename T>
   array(
       std::initializer_list<T> data,
-      std::vector<int> shape,
+      const std::vector<int>& shape,
       Dtype dtype = TypeToDtype<T>());
 
   /* Build an array from a buffer */
   array(
       allocator::Buffer data,
-      std::vector<int> shape,
+      const std::vector<int>& shape,
       Dtype dtype,
       deleter_t deleter = allocator::free);
 
@@ -73,32 +71,32 @@ class array {
       this->array_desc_ = other.array_desc_;
     }
     return *this;
-  }
+  };
 
   /** The size of the array's datatype in bytes. */
   size_t itemsize() const {
     return size_of(dtype());
-  }
+  };
 
   /** The number of elements in the array. */
   size_t size() const {
     return array_desc_->size;
-  }
+  };
 
   /** The number of bytes in the array. */
   size_t nbytes() const {
     return size() * itemsize();
-  }
+  };
 
   /** The number of dimensions of the array. */
   size_t ndim() const {
     return array_desc_->shape.size();
-  }
+  };
 
   /** The shape of the array as a vector of integers. */
   const std::vector<int>& shape() const {
     return array_desc_->shape;
-  }
+  };
 
   /**
    *  Get the size of the corresponding dimension.
@@ -107,26 +105,17 @@ class array {
    *  bounds checking. */
   int shape(int dim) const {
     return shape().at(dim < 0 ? dim + ndim() : dim);
-  }
+  };
 
   /** The strides of the array. */
   const std::vector<size_t>& strides() const {
     return array_desc_->strides;
-  }
-
-  /**
-   *  Get the stride of the corresponding dimension.
-   *
-   *  This function supports negative indexing and provides
-   *  bounds checking. */
-  size_t strides(int dim) const {
-    return strides().at(dim < 0 ? dim + ndim() : dim);
-  }
+  };
 
   /** Get the arrays data type. */
   Dtype dtype() const {
     return array_desc_->dtype;
-  }
+  };
 
   /** Evaluate the array. */
   void eval();
@@ -160,10 +149,10 @@ class array {
 
     friend bool operator==(const ArrayIterator& a, const ArrayIterator& b) {
       return a.arr.id() == b.arr.id() && a.idx == b.idx;
-    }
+    };
     friend bool operator!=(const ArrayIterator& a, const ArrayIterator& b) {
       return !(a == b);
-    }
+    };
 
    private:
     const array& arr;
@@ -183,16 +172,22 @@ 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(
-      std::vector<std::vector<int>> shapes,
+      const std::vector<std::vector<int>>& shapes,
       const std::vector<Dtype>& dtypes,
-      const std::shared_ptr<Primitive>& primitive,
+      std::shared_ptr<Primitive> primitive,
       const std::vector<array>& inputs);
 
   /** A unique identifier for an array. */
@@ -209,7 +204,7 @@ class array {
     allocator::Buffer buffer;
     deleter_t d;
     Data(allocator::Buffer buffer, deleter_t d = allocator::free)
-        : buffer(buffer), d(d) {}
+        : buffer(buffer), d(d){};
     // Not copyable
     Data(const Data& d) = delete;
     Data& operator=(const Data& d) = delete;
@@ -230,22 +225,22 @@ class array {
   /** The array's primitive. */
   Primitive& primitive() const {
     return *(array_desc_->primitive);
-  }
+  };
 
   /** A shared pointer to the array's primitive. */
   std::shared_ptr<Primitive>& primitive_ptr() const {
     return array_desc_->primitive;
-  }
+  };
 
   /** Check if the array has an attached primitive or is a leaf node. */
   bool has_primitive() const {
     return array_desc_->primitive != nullptr;
-  }
+  };
 
   /** The array's inputs. */
   const std::vector<array>& inputs() const {
     return array_desc_->inputs;
-  }
+  };
 
   std::vector<array>& inputs() {
     return array_desc_->inputs;
@@ -259,12 +254,7 @@ class array {
   /** The array's siblings. */
   const std::vector<array>& siblings() const {
     return array_desc_->siblings;
-  }
-
-  /** The array's siblings. */
-  std::vector<array>& siblings() {
-    return array_desc_->siblings;
-  }
+  };
 
   void set_siblings(std::vector<array> siblings, uint16_t position) {
     array_desc_->siblings = std::move(siblings);
@@ -281,6 +271,11 @@ class array {
     outputs.push_back(*this);
     outputs.insert(outputs.end(), siblings().begin() + idx, siblings().end());
     return outputs;
+  };
+
+  /** The depth of the array in the graph. Evaluated arrays have depth 0. */
+  uint16_t graph_depth() const {
+    return array_desc_->depth;
   }
 
   /** Detach the array from the graph. */
@@ -289,19 +284,19 @@ class array {
   /** Get the Flags bit-field. */
   const Flags& flags() const {
     return array_desc_->flags;
-  }
+  };
 
   /** The size (in elements) of the underlying buffer the array points to. */
   size_t data_size() const {
     return array_desc_->data_size;
-  }
+  };
 
   allocator::Buffer& buffer() {
     return array_desc_->data->buffer;
-  }
+  };
   const allocator::Buffer& buffer() const {
     return array_desc_->data->buffer;
-  }
+  };
 
   // Return a copy of the shared pointer
   // to the array::Data struct
@@ -312,35 +307,16 @@ class array {
   template <typename T>
   T* data() {
     return static_cast<T*>(array_desc_->data_ptr);
-  }
+  };
 
   template <typename T>
   const T* data() const {
     return static_cast<T*>(array_desc_->data_ptr);
-  }
+  };
 
-  enum Status { unscheduled, scheduled, available };
-
-  bool is_available() const {
-    return status() == Status::available;
-  }
-
-  Status status() const {
-    return array_desc_->status;
-  }
-
-  void set_status(Status s) const {
-    array_desc_->status = s;
-  }
-
-  // Get the array's shared event
-  Event& event() const {
-    return array_desc_->event;
-  }
-
-  // Attach an event to a not yet evaluated array
-  void attach_event(Event e) const {
-    array_desc_->event = std::move(e);
+  // Check if the array has been evaluated
+  bool is_evaled() const {
+    return array_desc_->data != nullptr;
   }
 
   // Mark the array as a tracer array (true) or not.
@@ -368,21 +344,12 @@ class array {
 
   void copy_shared_buffer(const array& other);
 
-  void move_shared_buffer(
-      array other,
-      const std::vector<size_t>& strides,
-      Flags flags,
-      size_t data_size,
-      size_t offset = 0);
-
   void move_shared_buffer(array other);
 
   void overwrite_descriptor(const array& other) {
     array_desc_ = other.array_desc_;
   }
 
-  ~array();
-
  private:
   // Initialize the arrays data
   template <typename It>
@@ -393,12 +360,7 @@ class array {
     std::vector<size_t> strides;
     size_t size;
     Dtype dtype;
-    std::shared_ptr<Primitive> primitive;
-
-    Status status;
-
-    // An event on the array used for synchronization
-    Event event;
+    std::shared_ptr<Primitive> primitive{nullptr};
 
     // Indicates an array is being used in a graph transform
     // and should not be detached from the graph
@@ -406,7 +368,7 @@ class array {
 
     // This is a shared pointer so that *different* arrays
     // can share the underlying data buffer.
-    std::shared_ptr<Data> data;
+    std::shared_ptr<Data> data{nullptr};
 
     // Properly offset data pointer
     void* data_ptr{nullptr};
@@ -426,26 +388,29 @@ class array {
     // The arrays position in the output list
     uint32_t position{0};
 
-    explicit ArrayDesc(std::vector<int> shape, Dtype dtype);
+    // The depth of the array in the graph.
+    uint16_t depth{0};
+
+    explicit ArrayDesc(const std::vector<int>& shape, Dtype dtype);
 
     explicit ArrayDesc(
-        std::vector<int> shape,
+        const std::vector<int>& shape,
         Dtype dtype,
         std::shared_ptr<Primitive> primitive,
-        std::vector<array> inputs);
+        const std::vector<array>& inputs);
 
-    ~ArrayDesc();
-
-   private:
-    // Initialize size, strides, and other metadata
-    void init();
+    explicit ArrayDesc(
+        std::vector<int>&& shape,
+        Dtype dtype,
+        std::shared_ptr<Primitive> primitive,
+        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_;
+  std::shared_ptr<ArrayDesc> array_desc_{nullptr};
 };
 
 template <typename T>
@@ -457,9 +422,9 @@ array::array(T val, Dtype dtype /* = TypeToDtype<T>() */)
 template <typename It>
 array::array(
   It data,
-  std::vector<int> shape,
+  const std::vector<int>& shape,
   Dtype dtype /* = TypeToDtype<typename std::iterator_traits<It>::value_type>() */) :
-    array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+    array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
   init(data);
 }
 
@@ -476,9 +441,9 @@ array::array(
 template <typename T>
 array::array(
     std::initializer_list<T> data,
-    std::vector<int> shape,
+    const std::vector<int>& shape,
     Dtype dtype /* = TypeToDtype<T>() */)
-    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(shape, dtype)) {
   if (data.size() != size()) {
     throw std::invalid_argument(
         "Data size and provided shape mismatch in array construction.");
@@ -500,11 +465,10 @@ T array::item() const {
   if (size() != 1) {
     throw std::invalid_argument("item can only be called on arrays of size 1.");
   }
-  if (status() == Status::unscheduled) {
+  if (!is_evaled()) {
     throw std::invalid_argument(
         "item() const can only be called on evaled arrays");
   }
-  const_cast<array*>(this)->eval();
   return *data<T>();
 }
 
@@ -554,15 +518,4 @@ void array::init(It src) {
   }
 }
 
-/* Utilities for determining whether a template parameter is array. */
-template <typename T>
-inline constexpr bool is_array_v =
-    std::is_same_v<std::remove_cv_t<std::remove_reference_t<T>>, array>;
-
-template <typename... T>
-inline constexpr bool is_arrays_v = (is_array_v<T> && ...);
-
-template <typename... T>
-using enable_for_arrays_t = typename std::enable_if_t<is_arrays_v<T...>>;
-
 } // namespace mlx::core

mlx/backend/accelerate/conv.cpp

@@ -1,9 +1,9 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <cassert>
 
-#include <Accelerate/Accelerate.h>
 #include <simd/vector.h>
+#include <vecLib/vDSP.h>
 
 #include "mlx/backend/common/copy.h"
 #include "mlx/primitives.h"

mlx/backend/accelerate/matmul.cpp

@@ -1,8 +1,9 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <cassert>
 
-#include <Accelerate/Accelerate.h>
+#include <vecLib/BNNS/bnns.h>
+#include <vecLib/cblas_new.h>
 
 #include "mlx/backend/accelerate/utils.h"
 #include "mlx/backend/common/copy.h"
@@ -195,40 +196,6 @@ inline void matmul_bnns(const array& a_pre, const array& b_pre, array& out) {
   return matmul_bnns_general(a_pre, b_pre, out);
 }
 
-template <typename T>
-inline void mask_matrix(
-    T* data,
-    const bool* mask,
-    int tile_size,
-    const int X,
-    const int Y,
-    const size_t X_data_str,
-    const size_t Y_data_str,
-    const size_t X_mask_str,
-    const size_t Y_mask_str) {
-  int tX = (X + tile_size - 1) / tile_size;
-  int tY = (Y + tile_size - 1) / tile_size;
-
-  for (int i = 0; i < tX; i++) {
-    for (int j = 0; j < tY; j++) {
-      bool do_mask = mask[i * X_mask_str + j * Y_mask_str];
-      if (!do_mask) {
-        int loc_x = i * tile_size;
-        int loc_y = j * tile_size;
-        T* data_block = data + loc_x * X_data_str + loc_y * Y_data_str;
-
-        int size_x = std::min(tile_size, X - loc_x);
-        int size_y = std::min(tile_size, Y - loc_y);
-        for (int ii = 0; ii < size_x; ii++) {
-          for (int jj = 0; jj < size_y; jj++) {
-            data_block[ii * X_data_str + jj * Y_data_str] = T(0.);
-          }
-        }
-      }
-    }
-  }
-}
-
 } // namespace
 
 void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {

mlx/backend/accelerate/primitives.cpp

@@ -3,7 +3,8 @@
 #include <cassert>
 #include <cmath>
 
-#include <Accelerate/Accelerate.h>
+#include <vecLib/vDSP.h>
+#include <vecLib/vForce.h>
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/binary.h"
@@ -30,27 +31,21 @@ DEFAULT(ArgPartition)
 DEFAULT(ArgReduce)
 DEFAULT(ArgSort)
 DEFAULT(AsStrided)
-DEFAULT(BlockMaskedMM)
 DEFAULT(Broadcast)
 DEFAULT(Ceil)
 DEFAULT(Concatenate)
-DEFAULT(Conjugate)
 DEFAULT(Copy)
-DEFAULT_MULTI(CustomTransforms)
+DEFAULT_MULTI(CustomVJP)
 DEFAULT_MULTI(Depends)
 DEFAULT_MULTI(DivMod)
-DEFAULT(NumberOfElements)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
 DEFAULT(FFT)
 DEFAULT(Floor)
 DEFAULT(Gather)
-DEFAULT(GatherMM)
-DEFAULT(GatherQMM)
 DEFAULT(Greater)
 DEFAULT(GreaterEqual)
-DEFAULT(Hadamard)
 DEFAULT(Less)
 DEFAULT(LessEqual)
 DEFAULT(Load)
@@ -73,14 +68,10 @@ DEFAULT(Select)
 DEFAULT(Sigmoid)
 DEFAULT(Sign)
 DEFAULT(Slice)
-DEFAULT(SliceUpdate)
 DEFAULT_MULTI(Split)
 DEFAULT(Sort)
 DEFAULT(StopGradient)
-DEFAULT_MULTI(SVD)
 DEFAULT(Transpose)
-DEFAULT(Inverse)
-DEFAULT(Cholesky)
 
 void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
@@ -102,7 +93,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary_op<float>(
+    binary(
         a,
         b,
         out,
@@ -117,7 +108,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vadd((const float*)a, 1, (const float*)b, 1, (float*)o, 1, n);
         });
   } else if (a.dtype() == int32) {
-    binary_op<int>(
+    binary(
         a,
         b,
         out,
@@ -132,7 +123,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vaddi((const int*)a, 1, (const int*)b, 1, (int*)o, 1, n);
         });
   } else {
-    eval(inputs, out);
+    binary(a, b, out, [](auto x, auto y) { return x + y; });
   }
 }
 
@@ -196,26 +187,6 @@ void ArcTan::eval_cpu(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void ArcTan2::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  auto& a = inputs[0];
-  auto& b = inputs[1];
-  if (out.dtype() == float32 && a.flags().row_contiguous &&
-      b.flags().row_contiguous) {
-    if (a.is_donatable()) {
-      out.copy_shared_buffer(a);
-    } else if (b.is_donatable()) {
-      out.copy_shared_buffer(b);
-    } else {
-      out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    }
-    int size = a.data_size();
-    vvatan2f(out.data<float>(), a.data<float>(), b.data<float>(), &size);
-  } else {
-    eval(inputs, out);
-  }
-}
-
 void ArcTanh::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
@@ -287,7 +258,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == int32) {
-    binary_op<int>(
+    binary(
         a,
         b,
         out,
@@ -300,7 +271,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vdivi((const int*)b, 1, (const int*)a, 1, (int*)o, 1, n);
         });
   } else if (a.dtype() == float32) {
-    binary_op<float>(
+    binary(
         a,
         b,
         out,
@@ -315,7 +286,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vdiv((const float*)b, 1, (const float*)a, 1, (float*)o, 1, n);
         });
   } else {
-    eval(inputs, out);
+    binary(a, b, out, [](auto x, auto y) { return x / y; });
   }
 }
 
@@ -326,21 +297,12 @@ void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
     set_unary_output_data(in, out);
     auto size = in.data_size();
     vvexpf(out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
+  } else if (is_floating_point(out.dtype())) {
+    unary_fp(in, out, [](auto x) { return std::exp(x); });
   } else {
-    eval(inputs, out);
-  }
-}
-
-void Expm1::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  const auto& in = inputs[0];
-  if (out.dtype() == float32 && in.flags().contiguous) {
-    set_unary_output_data(in, out);
-    auto size = in.data_size();
-    vvexpm1f(
-        out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
-  } else {
-    eval(inputs, out);
+    throw std::invalid_argument(
+        "[exp] Cannot exponentiate elements in array"
+        " with non floating point type.");
   }
 }
 
@@ -389,8 +351,12 @@ void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
     auto size = in.data_size();
     vvlog1pf(
         out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
+  } else if (is_floating_point(out.dtype())) {
+    unary_fp(in, out, [](auto x) { return std::log1p(x); });
   } else {
-    eval(inputs, out);
+    throw std::invalid_argument(
+        "[log1p] Cannot compute log of elements in array with"
+        " non floating point type.");
   }
 }
 
@@ -400,7 +366,7 @@ void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary_op<float>(
+    binary(
         a,
         b,
         out,
@@ -415,7 +381,7 @@ void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vmul((const float*)a, 1, (const float*)b, 1, (float*)o, 1, n);
         });
   } else {
-    eval(inputs, out);
+    binary(a, b, out, [](auto x, auto y) { return x * y; });
   }
 }
 
@@ -426,7 +392,7 @@ void Negative::eval_cpu(const std::vector<array>& inputs, array& out) {
     set_unary_output_data(in, out);
     vDSP_vneg(in.data<float>(), 1, out.data<float>(), 1, in.data_size());
   } else {
-    eval(inputs, out);
+    unary(in, out, [](auto x) { return -x; });
   }
 }
 
@@ -513,7 +479,7 @@ void Square::eval_cpu(const std::vector<array>& inputs, array& out) {
     auto size = in.data_size();
     vDSP_vsq(in.data<float>(), 1, out.data<float>(), 1, size);
   } else {
-    eval(inputs, out);
+    unary(in, out, [](auto x) { return x * x; });
   }
 }
 
@@ -539,7 +505,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary_op<float>(
+    binary(
         a,
         b,
         out,
@@ -557,7 +523,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vsub((const float*)b, 1, (const float*)a, 1, (float*)o, 1, n);
         });
   } else if (a.dtype() == int32) {
-    binary_op<int>(
+    binary(
         a,
         b,
         out,
@@ -569,7 +535,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
         },
         UseDefaultBinaryOp());
   } else {
-    eval(inputs, out);
+    binary(a, b, out, [](auto x, auto y) { return x - y; });
   }
 }
 

mlx/backend/accelerate/reduce.cpp

@@ -2,73 +2,86 @@
 
 #include <cassert>
 
-#include <Accelerate/Accelerate.h>
 #include <simd/vector.h>
+#include <vecLib/vDSP.h>
 
 #include "mlx/backend/common/reduce.h"
 #include "mlx/primitives.h"
 
 namespace mlx::core {
 
-namespace {
-
-template <typename T, typename VT>
-struct MinReduction {
-  T operator()(const T& a, const T& b) {
-    return std::min(a, b);
-  }
-
-  VT operator()(VT a, VT b) {
-    return simd_min(a, b);
-  }
-};
-
-template <typename T, typename VT>
-struct MaxReduction {
-  T operator()(const T& a, const T& b) {
-    return std::max(a, b);
-  }
-
-  VT operator()(VT a, VT b) {
-    return simd_max(a, b);
-  }
-};
-
-template <typename T, typename VT>
-struct SumReduction {
-  T operator()(const T& a, const T& b) {
-    return a + b;
-  }
-
-  VT operator()(VT a, VT b) {
-    return a + b;
-  }
-};
-
-template <typename T, typename VT, int N, typename Reduction>
-struct StridedReduce {
-  void operator()(const T* x, T* accum, int size, size_t stride) {
-    Reduction op;
-
-    for (int i = 0; i < size; i++) {
-      size_t s = stride;
-      T* a = accum;
-      while (s >= N) {
-        *(VT*)a = op((*(VT*)x), (*(VT*)a));
-        x += N;
-        a += N;
-        s -= N;
-      }
-      while (s-- > 0) {
-        *a = op(*a, *x);
-        a++;
-        x++;
-      }
+template <typename T, typename VT, int N>
+void _vectorized_strided_sum(const T* x, T* accum, int size, size_t stride) {
+  for (int i = 0; i < size; i++) {
+    size_t s = stride;
+    T* a = accum;
+    while (s >= N) {
+      VT val = (*(VT*)x);
+      *(VT*)a += val;
+      x += N;
+      a += N;
+      s -= N;
+    }
+    while (s-- > 0) {
+      *a++ += *x++;
     }
   }
-};
+}
 
-} // namespace
+// TODO: Add proper templates for the strided reduce algorithm so we don't have
+// to write max/min/sum etc.
+template <typename T, typename VT, int N>
+void _vectorized_strided_max(const T* x, T* accum, int size, size_t stride) {
+  for (int i = 0; i < size; i++) {
+    size_t s = stride;
+    T* a = accum;
+    while (s >= N) {
+      *(VT*)a = simd_max((*(VT*)x), (*(VT*)a));
+      x += N;
+      a += N;
+      s -= N;
+    }
+    while (s-- > 0) {
+      *a = std::max(*a, *x);
+      a++;
+      x++;
+    }
+  }
+}
+
+template <typename T, typename VT, int N>
+void _vectorized_strided_min(const T* x, T* accum, int size, size_t stride) {
+  for (int i = 0; i < size; i++) {
+    size_t s = stride;
+    T* a = accum;
+    while (s >= N) {
+      *(VT*)a = simd_min((*(VT*)x), (*(VT*)a));
+      x += N;
+      a += N;
+      s -= N;
+    }
+    while (s-- > 0) {
+      *a = std::min(*a, *x);
+      a++;
+      x++;
+    }
+  }
+}
+
+template <typename T, typename VT, int N>
+void _vectorized_sum(const T* x, T* accum, int size) {
+  VT _sum = {0};
+  while (size >= N) {
+    _sum += (*(VT*)x);
+    x += N;
+    size -= N;
+  }
+  T sum = _sum[0];
+  for (int i = 1; i < N; i++) {
+    sum += _sum[i];
+  }
+  *accum += sum;
+}
 
 void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
@@ -81,11 +94,10 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
           out,
           axes_,
           0,
-          StridedReduce<
-              float,
-              simd_float16,
-              16,
-              SumReduction<float, simd_float16>>(),
+          [](const auto* x, auto* accum, int size, size_t stride) {
+            _vectorized_strided_sum<float, simd_float16, 16>(
+                (const float*)x, (float*)accum, size, stride);
+          },
           [](const auto* x, auto* accum, int size) {
             float acc;
             vDSP_sve((const float*)x, 1, &acc, size);
@@ -99,11 +111,10 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
           out,
           axes_,
           -std::numeric_limits<float>::infinity(),
-          StridedReduce<
-              float,
-              simd_float16,
-              16,
-              MaxReduction<float, simd_float16>>(),
+          [](const auto* x, auto* accum, int size, size_t stride) {
+            _vectorized_strided_max<float, simd_float16, 16>(
+                (const float*)x, (float*)accum, size, stride);
+          },
           [](const auto* x, auto* accum, int size) {
             float max;
             vDSP_maxv((const float*)x, 1, &max, size);
@@ -117,11 +128,10 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
           out,
           axes_,
           std::numeric_limits<float>::infinity(),
-          StridedReduce<
-              float,
-              simd_float16,
-              16,
-              MinReduction<float, simd_float16>>(),
+          [](const auto* x, auto* accum, int size, size_t stride) {
+            _vectorized_strided_min<float, simd_float16, 16>(
+                (const float*)x, (float*)accum, size, stride);
+          },
           [](const auto* x, auto* accum, int size) {
             float min;
             vDSP_minv((const float*)x, 1, &min, size);

mlx/backend/accelerate/softmax.cpp

@@ -1,12 +1,9 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <cassert>
 #include <limits>
 
-#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
 #include <arm_neon.h>
-#endif
-
 #include <simd/math.h>
 #include <simd/vector.h>
 
@@ -56,26 +53,25 @@ inline simd_float16 simd_fast_exp(simd_float16 x) {
   return (*(simd_float16*)&epart) * x;
 }
 
-#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
 /**
  * The ARM neon equivalent of the fast exp above.
  */
 inline float16x8_t neon_fast_exp(float16x8_t x) {
-  x = vmulq_f16(x, vdupq_n_f16(float16_t(1.442695f))); // multiply with log_2(e)
-  x = vmaxq_f16(x, vdupq_n_f16(float16_t(-14.f))); // clamp under with -14
-  x = vminq_f16(x, vdupq_n_f16(float16_t(14.f))); // clamp over with 14
+  x = vmulq_f16(x, vdupq_n_f16(1.442695)); // multiply with log_2(e)
+  x = vmaxq_f16(x, vdupq_n_f16(-14)); // clamp under with -14
+  x = vminq_f16(x, vdupq_n_f16(14)); // clamp over with 14
 
-  float16x8_t ipart = vrndmq_f16(vaddq_f16(x, vdupq_n_f16(float16_t(0.5f))));
+  float16x8_t ipart = vrndmq_f16(vaddq_f16(x, vdupq_n_f16(0.5)));
   float16x8_t fpart = vsubq_f16(x, ipart);
 
-  x = vdupq_n_f16(float16_t(1.535336188319500e-4f));
-  x = vfmaq_f16(vdupq_n_f16(float16_t(1.339887440266574e-3f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(1.339887440266574e-3f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(9.618437357674640e-3f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(5.550332471162809e-2f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(2.402264791363012e-1f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(6.931472028550421e-1f)), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(float16_t(1.000000000000000f)), x, fpart);
+  x = vdupq_n_f16(1.535336188319500e-4f);
+  x = vfmaq_f16(vdupq_n_f16(1.339887440266574e-3f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(1.339887440266574e-3f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(9.618437357674640e-3f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(5.550332471162809e-2f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(2.402264791363012e-1f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(6.931472028550421e-1f), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(1.000000000000000f), x, fpart);
 
   // generate 2**ipart in the floating point representation using integer
   // bitshifting
@@ -111,55 +107,6 @@ inline float16_t neon_reduce_add(float16x8_t x) {
   return vget_lane_f16(y, 0);
 }
 
-template <typename T, typename VT>
-struct NeonFp16SimdOps {
-  VT init(T a) {
-    return vdupq_n_f16(a);
-  }
-
-  VT load(const T* a) {
-    return vld1q_f16(a);
-  }
-
-  void store(T* dst, VT x) {
-    vst1q_f16(dst, x);
-  }
-
-  VT max(VT a, VT b) {
-    return vmaxq_f16(a, b);
-  }
-
-  VT exp(VT x) {
-    return neon_fast_exp(x);
-  }
-
-  VT add(VT a, VT b) {
-    return vaddq_f16(a, b);
-  }
-
-  VT sub(VT a, T b) {
-    return vsubq_f16(a, vdupq_n_f16(b));
-  }
-
-  VT mul(VT a, VT b) {
-    return vmulq_f16(a, b);
-  }
-
-  VT mul(VT a, T b) {
-    return vmulq_f16(a, vdupq_n_f16(b));
-  }
-
-  T reduce_max(VT x) {
-    return neon_reduce_max(x);
-  }
-
-  T reduce_add(VT x) {
-    return neon_reduce_add(x);
-  }
-};
-
-#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-
 template <typename T, typename VT>
 struct AccelerateSimdOps {
   VT init(T a) {
@@ -176,7 +123,7 @@ struct AccelerateSimdOps {
 
   VT max(VT a, VT b) {
     return simd_max(a, b);
-  }
+  };
 
   VT exp(VT x) {
     return simd_fast_exp(x);
@@ -207,7 +154,54 @@ struct AccelerateSimdOps {
   }
 };
 
-template <typename T, typename AccT, typename VT, typename Ops, int N>
+template <typename T, typename VT>
+struct NeonFp16SimdOps {
+  VT init(T a) {
+    return vdupq_n_f16(a);
+  }
+
+  VT load(const T* a) {
+    return vld1q_f16(a);
+  }
+
+  void store(T* dst, VT x) {
+    vst1q_f16(dst, x);
+  }
+
+  VT max(VT a, VT b) {
+    return vmaxq_f16(a, b);
+  };
+
+  VT exp(VT x) {
+    return neon_fast_exp(x);
+  }
+
+  VT add(VT a, VT b) {
+    return vaddq_f16(a, b);
+  }
+
+  VT sub(VT a, T b) {
+    return vsubq_f16(a, vdupq_n_f16(b));
+  }
+
+  VT mul(VT a, VT b) {
+    return vmulq_f16(a, b);
+  }
+
+  VT mul(VT a, T b) {
+    return vmulq_f16(a, vdupq_n_f16(b));
+  }
+
+  T reduce_max(VT x) {
+    return neon_reduce_max(x);
+  }
+
+  T reduce_add(VT x) {
+    return neon_reduce_add(x);
+  }
+};
+
+template <typename T, typename VT, typename Ops, int N>
 void softmax(const array& in, array& out) {
   Ops ops;
 
@@ -224,21 +218,13 @@ void softmax(const array& in, array& out) {
     VT vmaximum = ops.init(-std::numeric_limits<float>::infinity());
     size_t s = M;
     while (s >= N) {
-      VT vals;
-      if constexpr (std::is_same<T, AccT>::value) {
-        vals = ops.load(current_in_ptr);
-      } else {
-        for (int i = 0; i < N; ++i) {
-          vals[i] = static_cast<AccT>(current_in_ptr[i]);
-        }
-      }
-      vmaximum = ops.max(vals, vmaximum);
+      vmaximum = ops.max(ops.load(current_in_ptr), vmaximum);
       current_in_ptr += N;
       s -= N;
     }
-    AccT maximum = ops.reduce_max(vmaximum);
+    T maximum = ops.reduce_max(vmaximum);
     while (s-- > 0) {
-      maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
+      maximum = std::max(maximum, *current_in_ptr);
       current_in_ptr++;
     }
 
@@ -248,29 +234,18 @@ void softmax(const array& in, array& out) {
     current_in_ptr = in_ptr;
     s = M;
     while (s >= N) {
-      VT vexp;
-      if constexpr (std::is_same<T, AccT>::value) {
-        vexp = ops.load(current_in_ptr);
-      } else {
-        for (int i = 0; i < N; ++i) {
-          vexp[i] = static_cast<AccT>(current_in_ptr[i]);
-        }
-      }
-      vexp = ops.exp(ops.sub(vexp, maximum));
-      if constexpr (std::is_same<T, AccT>::value) {
-        ops.store(current_out_ptr, vexp);
-      }
+      VT vexp = ops.exp(ops.sub(*(VT*)current_in_ptr, maximum));
+      ops.store(current_out_ptr, vexp);
+      *(VT*)current_out_ptr = vexp;
       vnormalizer = ops.add(vnormalizer, vexp);
       current_in_ptr += N;
       current_out_ptr += N;
       s -= N;
     }
-    AccT normalizer = ops.reduce_add(vnormalizer);
+    T normalizer = ops.reduce_add(vnormalizer);
     while (s-- > 0) {
-      AccT _exp = std::exp(*current_in_ptr - maximum);
-      if (std::is_same<T, AccT>::value) {
-        *current_out_ptr = _exp;
-      }
+      T _exp = std::exp(*current_in_ptr - maximum);
+      *current_out_ptr = _exp;
       normalizer += _exp;
       current_in_ptr++;
       current_out_ptr++;
@@ -279,33 +254,14 @@ void softmax(const array& in, array& out) {
 
     // Normalize
     current_out_ptr = out_ptr;
-    current_in_ptr = in_ptr;
     s = M;
     while (s >= N) {
-      if constexpr (std::is_same<T, AccT>::value) {
-        ops.store(current_out_ptr, ops.mul(*(VT*)current_out_ptr, normalizer));
-      } else {
-        VT vexp;
-        for (int i = 0; i < N; ++i) {
-          vexp[i] = static_cast<AccT>(current_in_ptr[i]);
-        }
-        vexp = ops.mul(ops.exp(ops.sub(vexp, maximum)), normalizer);
-        for (int i = 0; i < N; ++i) {
-          current_out_ptr[i] = vexp[i];
-        }
-        current_in_ptr += N;
-      }
+      ops.store(current_out_ptr, ops.mul(*(VT*)current_out_ptr, normalizer));
       current_out_ptr += N;
       s -= N;
     }
     while (s-- > 0) {
-      if constexpr (std::is_same<T, AccT>::value) {
-        *current_out_ptr *= normalizer;
-      } else {
-        AccT _exp = std::exp(*current_in_ptr - maximum);
-        *current_out_ptr = static_cast<T>(_exp * normalizer);
-        current_in_ptr++;
-      }
+      *current_out_ptr *= normalizer;
       current_out_ptr++;
     }
   }
@@ -352,33 +308,15 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
           "Softmax is defined only for floating point types");
       break;
     case float32:
-      softmax<
-          float,
-          float,
-          simd_float16,
-          AccelerateSimdOps<float, simd_float16>,
-          16>(in, out);
+      softmax<float, simd_float16, AccelerateSimdOps<float, simd_float16>, 16>(
+          in, out);
       break;
     case float16:
-      if (precise_) {
-        softmax<
-            float16_t,
-            float,
-            simd_float16,
-            AccelerateSimdOps<float, simd_float16>,
-            16>(in, out);
-      } else {
-#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-        softmax<
-            float16_t,
-            float16_t,
-            float16x8_t,
-            NeonFp16SimdOps<float16_t, float16x8_t>,
-            8>(in, out);
-#else // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-        eval(inputs, out); // Redirect to common backend for consistency
-#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-      }
+      softmax<
+          float16_t,
+          float16x8_t,
+          NeonFp16SimdOps<float16_t, float16x8_t>,
+          8>(in, out);
       break;
     case bfloat16:
       eval(inputs, out);

mlx/backend/accelerate/utils.h

@@ -1,8 +1,8 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #pragma once
 
-#include <Accelerate/Accelerate.h>
+#include <vecLib/BNNS/bnns.h>
 #include "mlx/dtype.h"
 
 namespace mlx::core {

mlx/backend/common/CMakeLists.txt

@@ -37,29 +37,22 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/erf.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/masked_mm.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/reduce_utils.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
   ${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_SOURCE_DIR}/cholesky.cpp
   ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
 )
 

mlx/backend/common/binary.cpp

@@ -179,16 +179,18 @@ void LogAddExp::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  if (out.dtype() == float32) {
-    binary_op<float>(a, b, out, detail::LogAddExp());
-  } else if (out.dtype() == float16) {
-    binary_op<float16_t>(a, b, out, detail::LogAddExp());
-  } else if (out.dtype() == bfloat16) {
-    binary_op<bfloat16_t>(a, b, out, detail::LogAddExp());
-  } else if (issubdtype(out.dtype(), inexact)) {
-    std::ostringstream err;
-    err << "[logaddexp] Does not support " << out.dtype();
-    throw std::invalid_argument(err.str());
+  if (is_floating_point(out.dtype())) {
+    if (out.dtype() == float32) {
+      binary_op<float>(a, b, out, detail::LogAddExp());
+    } else if (out.dtype() == float16) {
+      binary_op<float16_t>(a, b, out, detail::LogAddExp());
+    } else if (out.dtype() == bfloat16) {
+      binary_op<bfloat16_t>(a, b, out, detail::LogAddExp());
+    } else {
+      std::ostringstream err;
+      err << "[logaddexp] Does not support " << out.dtype();
+      throw std::invalid_argument(err.str());
+    }
   } else {
     throw std::invalid_argument(
         "[logaddexp] Cannot compute logaddexp for arrays with"
@@ -196,20 +198,6 @@ void LogAddExp::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void LogicalAnd::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2); // LogicalAnd requires two input arrays
-  auto& in1 = inputs[0];
-  auto& in2 = inputs[1];
-  binary(in1, in2, out, detail::LogicalAnd());
-}
-
-void LogicalOr::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2); // LogicalOr requires two input arrays
-  auto& in1 = inputs[0];
-  auto& in2 = inputs[1];
-  binary(in1, in2, out, detail::LogicalOr());
-}
-
 void Maximum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
@@ -250,82 +238,4 @@ void Subtract::eval(const std::vector<array>& inputs, array& out) {
   binary(a, b, out, detail::Subtract());
 }
 
-void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  auto& a = inputs[0];
-  auto& b = inputs[1];
-  auto dispatch_type = [&a, &b, &out](auto op) {
-    switch (out.dtype()) {
-      case bool_:
-        binary_op<bool>(a, b, out, op);
-      case uint8:
-        binary_op<uint8_t>(a, b, out, op);
-        break;
-      case uint16:
-        binary_op<uint16_t>(a, b, out, op);
-        break;
-      case uint32:
-        binary_op<uint32_t>(a, b, out, op);
-        break;
-      case uint64:
-        binary_op<uint64_t>(a, b, out, op);
-        break;
-      case int8:
-        binary_op<int8_t>(a, b, out, op);
-        break;
-      case int16:
-        binary_op<int16_t>(a, b, out, op);
-        break;
-      case int32:
-        binary_op<int32_t>(a, b, out, op);
-        break;
-      case int64:
-        binary_op<int64_t>(a, b, out, op);
-        break;
-      default:
-        throw std::runtime_error(
-            "[BitwiseBinary::eval_cpu] Type not supported");
-        break;
-    }
-  };
-  switch (op_) {
-    case BitwiseBinary::And:
-      dispatch_type(detail::BitwiseAnd());
-      break;
-    case BitwiseBinary::Or:
-      dispatch_type(detail::BitwiseOr());
-      break;
-    case BitwiseBinary::Xor:
-      dispatch_type(detail::BitwiseXor());
-      break;
-    case BitwiseBinary::LeftShift:
-      dispatch_type(detail::LeftShift());
-      break;
-    case BitwiseBinary::RightShift:
-      dispatch_type(detail::RightShift());
-      break;
-  }
-}
-
-void ArcTan2::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  const auto& a = inputs[0];
-  const auto& b = inputs[1];
-  if (out.dtype() == float32) {
-    binary_op<float>(a, b, out, detail::ArcTan2());
-  } else if (out.dtype() == float16) {
-    binary_op<float16_t>(a, b, out, detail::ArcTan2());
-  } else if (out.dtype() == bfloat16) {
-    binary_op<bfloat16_t>(a, b, out, detail::ArcTan2());
-  } else if (issubdtype(out.dtype(), inexact)) {
-    std::ostringstream err;
-    err << "[arctan2] Does not support " << out.dtype();
-    throw std::invalid_argument(err.str());
-  } else {
-    throw std::invalid_argument(
-        "[arctan2] Cannot compute inverse tangent for arrays"
-        " with non floating point type.");
-  }
-}
-
 } // namespace mlx::core

mlx/backend/common/binary.h

@@ -1,8 +1,6 @@
 // Copyright © 2023 Apple Inc.
 
 #pragma once
-#include <cassert>
-
 #include "mlx/allocator.h"
 #include "mlx/array.h"
 #include "mlx/backend/common/utils.h"

mlx/backend/common/cholesky.cpp

@@ -1,101 +0,0 @@
-// 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 {
-
-namespace {
-
-// Delegate to the Cholesky factorization taking into account differences in
-// LAPACK implementations (basically how to pass the 'uplo' string to fortran).
-int spotrf_wrapper(char uplo, float* matrix, int N) {
-  int info;
-
-#ifdef LAPACK_FORTRAN_STRLEN_END
-  spotrf_(
-      /* uplo = */ &uplo,
-      /* n = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info,
-      /* uplo_len = */ static_cast<size_t>(1));
-#else
-  spotrf_(
-      /* uplo = */ &uplo,
-      /* n = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info);
-#endif
-
-  return info;
-}
-
-} // namespace
-
-void cholesky_impl(const array& a, array& factor, bool upper) {
-  // Lapack uses the column-major convention. We take advantage of the fact that
-  // the matrix should be symmetric:
-  //   (A)ᵀ = A
-  // and that a column-major lower triangular matrix is a row-major upper
-  // triangular matrix, so uplo is the opposite of what we would expect from
-  // upper
-
-  char uplo = (upper) ? 'L' : 'U';
-
-  // The decomposition is computed in place, so just copy the input to the
-  // output.
-  copy(
-      a,
-      factor,
-      a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
-
-  const int N = a.shape(-1);
-  const size_t num_matrices = a.size() / (N * N);
-
-  float* matrix = factor.data<float>();
-
-  for (int i = 0; i < num_matrices; i++) {
-    // Compute Cholesky factorization.
-    int info = spotrf_wrapper(uplo, matrix, N);
-
-    // TODO: We do nothing when the matrix is not positive semi-definite
-    // because throwing an error would result in a crash. If we figure out how
-    // to catch errors from the implementation we should throw.
-    if (info < 0) {
-      std::stringstream msg;
-      msg << "[cholesky] Cholesky decomposition failed with error code "
-          << info;
-      throw std::runtime_error(msg.str());
-    }
-
-    // Zero out the upper/lower triangle while advancing the pointer to the
-    // next matrix at the same time.
-    for (int row = 0; row < N; row++) {
-      if (upper) {
-        std::fill(matrix, matrix + row, 0);
-      } else {
-        std::fill(matrix + row + 1, matrix + N, 0);
-      }
-      matrix += N;
-    }
-  }
-}
-
-void Cholesky::eval(const std::vector<array>& inputs, array& output) {
-  if (inputs[0].dtype() != float32) {
-    throw std::runtime_error("[Cholesky::eval] only supports float32.");
-  }
-  cholesky_impl(inputs[0], output, upper_);
-}
-
-} // namespace mlx::core

mlx/backend/common/common.cpp

@@ -1,304 +0,0 @@
-// Copyright © 2024 Apple Inc.
-#include <cassert>
-
-#include "mlx/backend/common/utils.h"
-#include "mlx/primitives.h"
-
-namespace mlx::core {
-
-void AsStrided::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-
-  auto& in = inputs[0];
-
-  if (!in.flags().row_contiguous) {
-    // Just ensuring that inputs[0] came from the ops which would ensure the
-    // input is row contiguous.
-    throw std::runtime_error(
-        "AsStrided must be used with row contiguous arrays only.");
-  }
-
-  // Compute the flags given the shape and strides
-  bool row_contiguous = true, col_contiguous = true;
-  size_t r = 1, c = 1;
-  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
-    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
-    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
-    r *= shape_[i];
-    c *= shape_[j];
-  }
-  auto flags = in.flags();
-  // TODO: Compute the contiguous flag in a better way cause now we are
-  //       unnecessarily strict.
-  flags.contiguous = row_contiguous || col_contiguous;
-  flags.row_contiguous = row_contiguous;
-  flags.col_contiguous = col_contiguous;
-
-  // There is no easy way to compute the actual data size so we use out.size().
-  // The contiguous flag will almost certainly not be set so no code should
-  // rely on data_size anyway.
-  size_t data_size = out.size();
-
-  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
-}
-
-void Broadcast::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  const auto& in = inputs[0];
-  if (out.size() == 0) {
-    out.set_data(nullptr);
-    return;
-  }
-  std::vector<size_t> strides(out.ndim(), 0);
-  int diff = out.ndim() - in.ndim();
-  for (int i = in.ndim() - 1; i >= 0; --i) {
-    strides[i + diff] = (in.shape()[i] == 1) ? 0 : in.strides()[i];
-  }
-  auto flags = in.flags();
-  if (out.size() > in.size()) {
-    flags.row_contiguous = flags.col_contiguous = false;
-  }
-  out.copy_shared_buffer(in, strides, flags, in.data_size());
-}
-
-void Copy::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  out.copy_shared_buffer(inputs[0]);
-}
-
-void CustomTransforms::eval(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() > outputs.size());
-  for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
-       i++, j++) {
-    outputs[i].copy_shared_buffer(inputs[j]);
-  }
-}
-
-void Depends::eval(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() > outputs.size());
-  for (int i = 0; i < outputs.size(); i++) {
-    outputs[i].copy_shared_buffer(inputs[i]);
-  }
-}
-
-void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  double numel = 1;
-  for (auto ax : axes_) {
-    numel *= inputs[0].shape(ax);
-  }
-
-  if (inverted_) {
-    numel = 1.0 / numel;
-  }
-
-  switch (out.dtype()) {
-    case bool_:
-      *out.data<bool>() = static_cast<bool>(numel);
-      break;
-    case uint8:
-      *out.data<uint8_t>() = static_cast<uint8_t>(numel);
-      break;
-    case uint16:
-      *out.data<uint16_t>() = static_cast<uint16_t>(numel);
-      break;
-    case uint32:
-      *out.data<uint32_t>() = static_cast<uint32_t>(numel);
-      break;
-    case uint64:
-      *out.data<uint64_t>() = static_cast<uint64_t>(numel);
-      break;
-    case int8:
-      *out.data<int8_t>() = static_cast<int8_t>(numel);
-      break;
-    case int16:
-      *out.data<int16_t>() = static_cast<int16_t>(numel);
-      break;
-    case int32:
-      *out.data<int32_t>() = static_cast<int32_t>(numel);
-      break;
-    case int64:
-      *out.data<int64_t>() = static_cast<int64_t>(numel);
-      break;
-    case float16:
-      *out.data<float16_t>() = static_cast<float16_t>(numel);
-      break;
-    case float32:
-      *out.data<float>() = static_cast<float>(numel);
-      break;
-    case bfloat16:
-      *out.data<bfloat16_t>() = static_cast<bfloat16_t>(numel);
-      break;
-    case complex64:
-      *out.data<complex64_t>() = static_cast<complex64_t>(numel);
-      break;
-  }
-}
-
-std::pair<bool, std::vector<size_t>> Reshape::prepare_reshape(
-    const array& in,
-    const array& out) {
-  // Special case for empty arrays or row contiguous arrays
-  if (in.size() == 0 || in.flags().row_contiguous) {
-    return {false, out.strides()};
-  }
-
-  // Special case for scalars
-  if (in.ndim() == 0) {
-    std::vector<size_t> out_strides(out.ndim(), 0);
-    return {false, out_strides};
-  }
-
-  // Firstly let's collapse all the contiguous dimensions of the input
-  auto [shape, _strides] = collapse_contiguous_dims(in);
-  auto& strides = _strides[0];
-
-  // If shapes fit exactly in the contiguous dims then no copy is necessary so
-  // let's check.
-  std::vector<size_t> out_strides;
-  bool copy_necessary = false;
-  int j = 0;
-  for (int i = 0; i < out.ndim(); i++) {
-    int N = out.shape(i);
-    if (j < shape.size() && shape[j] % N == 0) {
-      shape[j] /= N;
-      out_strides.push_back(shape[j] * strides[j]);
-      j += (shape[j] == 1);
-    } else if (N == 1) {
-      // i > 0 because otherwise j < shape.size() && shape[j] % 1 == 0
-      out_strides.push_back(out_strides.back());
-    } else {
-      copy_necessary = true;
-      break;
-    }
-  }
-
-  return {copy_necessary, out_strides};
-}
-
-void Reshape::shared_buffer_reshape(
-    const array& in,
-    const std::vector<size_t>& out_strides,
-    array& out) {
-  auto flags = in.flags();
-  if (flags.row_contiguous) {
-    // For row contiguous reshapes:
-    // - Shallow copy the buffer
-    // - If reshaping into a vector (all singleton dimensions except one) it
-    //    becomes col contiguous again.
-    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
-    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
-  }
-  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
-}
-
-void Split::eval(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  assert(inputs.size() == 1);
-
-  auto& in = inputs[0];
-
-  auto compute_new_flags = [](const auto& shape,
-                              const auto& strides,
-                              size_t in_data_size,
-                              auto flags) {
-    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 = shape.size() - 1; ri >= 0; i++, ri--) {
-      flags.col_contiguous &= strides[i] == f_stride || shape[i] == 1;
-      flags.row_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
-      f_stride *= shape[i];
-      b_stride *= shape[ri];
-      if (strides[i] > 0) {
-        data_size *= shape[i];
-      }
-    }
-
-    if (data_size == 1) {
-      // Broadcasted scalar array is contiguous.
-      flags.contiguous = true;
-    } else if (data_size == in_data_size) {
-      // Means we sliced a broadcasted dimension so leave the "no holes" flag
-      // alone.
-    } else {
-      // We sliced something. So either we are row or col contiguous or we
-      // punched a hole.
-      flags.contiguous &= flags.row_contiguous || flags.col_contiguous;
-    }
-
-    return std::pair<decltype(flags), size_t>{flags, data_size};
-  };
-
-  std::vector<int> indices(1, 0);
-  indices.insert(indices.end(), indices_.begin(), indices_.end());
-  for (int i = 0; i < indices.size(); i++) {
-    size_t offset = indices[i] * in.strides()[axis_];
-    auto [new_flags, data_size] = compute_new_flags(
-        outputs[i].shape(), in.strides(), in.data_size(), in.flags());
-    outputs[i].copy_shared_buffer(
-        in, in.strides(), new_flags, data_size, offset);
-  }
-}
-
-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 StopGradient::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  out.copy_shared_buffer(inputs[0]);
-}
-
-void Transpose::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  std::vector<size_t> out_strides(out.ndim());
-  auto& in = inputs[0];
-  for (int ax = 0; ax < axes_.size(); ++ax) {
-    out_strides[ax] = in.strides()[axes_[ax]];
-  }
-
-  // Conditions for {row/col}_contiguous
-  // - array must be contiguous (no gaps)
-  // - underlying buffer size should have the same size as the array
-  // - cumulative product of shapes is equal to the strides (we can ignore axes
-  //   with size == 1)
-  //   - in the forward direction (column contiguous)
-  //   - in the reverse direction (row contiguous)
-  // - vectors are both row and col contiguous (hence if both row/col are
-  //   true, they stay true)
-  auto flags = in.flags();
-  if (flags.contiguous && in.data_size() == in.size()) {
-    size_t f_stride = 1;
-    size_t b_stride = 1;
-    flags.col_contiguous = true;
-    flags.row_contiguous = true;
-    for (int i = 0, ri = out.ndim() - 1; i < out.ndim(); ++i, --ri) {
-      flags.col_contiguous &= (out_strides[i] == f_stride || out.shape(i) == 1);
-      f_stride *= out.shape(i);
-      flags.row_contiguous &=
-          (out_strides[ri] == b_stride || out.shape(ri) == 1);
-      b_stride *= out.shape(ri);
-    }
-  }
-  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
-}
-
-} // namespace mlx::core

mlx/backend/common/compiled.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/common/compiled.h"
-#include "mlx/graph_utils.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
@@ -82,27 +81,13 @@ std::string build_lib_name(
     const std::vector<array>& outputs,
     const std::vector<array>& tape,
     const std::unordered_set<uintptr_t>& constant_ids) {
-  NodeNamer namer;
   std::ostringstream os;
   std::ostringstream constant_hasher;
 
-  // Fill the input names. This is not really necessary, I just like having A,
-  // B, C, ... as the inputs.
-  for (auto& x : inputs) {
-    namer.get_name(x);
-  }
-
   // The primitives describing the tape. For unary and binary primitives this
   // must be enough to describe the full computation.
   for (auto& a : tape) {
-    // name and type of output
-    os << namer.get_name(a) << kindof(a.dtype()) << a.itemsize();
-    // computation performed
     a.primitive().print(os);
-    // name of inputs to the function
-    for (auto& inp : a.inputs()) {
-      os << namer.get_name(inp);
-    }
   }
   os << "_";
 
@@ -126,102 +111,4 @@ std::string build_lib_name(
   return os.str();
 }
 
-bool compiled_check_contiguity(
-    const std::vector<array>& inputs,
-    const std::vector<int>& shape) {
-  bool contiguous = true;
-  bool all_contig = true;
-  bool all_row_contig = true;
-  bool all_col_contig = true;
-  int non_scalar_inputs = 0;
-  for (const auto& x : inputs) {
-    if (is_scalar(x)) {
-      continue;
-    }
-    non_scalar_inputs++;
-    bool shape_eq = x.shape() == shape;
-    all_contig &= (x.flags().contiguous && shape_eq);
-    all_row_contig &= (x.flags().row_contiguous && shape_eq);
-    all_col_contig &= (x.flags().col_contiguous && shape_eq);
-  }
-  if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) {
-    contiguous = false;
-  } else if (non_scalar_inputs == 1 && !all_contig) {
-    contiguous = false;
-  } else if (non_scalar_inputs == 0 && !shape.empty()) {
-    contiguous = false;
-  }
-  return contiguous;
-}
-
-void compiled_allocate_outputs(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs,
-    const std::vector<array>& inputs_,
-    const std::unordered_set<uintptr_t>& constant_ids_,
-    bool contiguous,
-    bool move_buffers /* = false */) {
-  if (contiguous) {
-    int o = 0;
-    std::vector<size_t> strides;
-    size_t data_size;
-    array::Flags flags;
-    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
-      auto& in = inputs[i];
-      // Conditions for donation
-      // - Correct size
-      // - Not a scalar
-      // - Donatable
-      // - Not a constant
-      if (in.itemsize() == outputs[o].itemsize() && !is_scalar(in) &&
-          in.is_donatable() &&
-          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
-        if (move_buffers) {
-          outputs[o++].move_shared_buffer(in);
-        } else {
-          outputs[o++].copy_shared_buffer(in);
-        }
-      }
-      // Get representative input flags to properly set non-donated outputs
-      if (strides.empty() && in.size() == outputs[0].size()) {
-        strides = in.strides();
-        flags = in.flags();
-        data_size = in.data_size();
-      }
-    }
-    for (; o < outputs.size(); ++o) {
-      outputs[o].set_data(
-          allocator::malloc_or_wait(data_size * outputs[o].itemsize()),
-          data_size,
-          strides,
-          flags);
-    }
-  } else {
-    int o = 0;
-    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
-      auto& in = inputs[i];
-      // Conditions for donation
-      // - Row contiguous
-      // - Donatable
-      // - Correct size
-      // - Not a constant
-      if (in.flags().row_contiguous && in.size() == outputs[o].size() &&
-          in.itemsize() == outputs[o].itemsize() && in.is_donatable() &&
-          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
-        if (move_buffers) {
-          outputs[o].move_shared_buffer(
-              in, outputs[o].strides(), in.flags(), in.data_size());
-        } else {
-          outputs[o].copy_shared_buffer(
-              in, outputs[o].strides(), in.flags(), in.data_size());
-        }
-        o++;
-      }
-    }
-    for (; o < outputs.size(); ++o) {
-      outputs[o].set_data(allocator::malloc_or_wait(outputs[o].nbytes()));
-    }
-  }
-}
-
 } // namespace mlx::core

mlx/backend/common/compiled.h

@@ -53,18 +53,4 @@ inline bool is_scalar(const array& x) {
   return x.ndim() == 0;
 }
 
-// Check if we can use a contiguous operation given inputs and the output shape
-bool compiled_check_contiguity(
-    const std::vector<array>& inputs,
-    const std::vector<int>& shape);
-
-// Allocate space for the outputs possibly with input donation
-void compiled_allocate_outputs(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs,
-    const std::vector<array>& inputs_,
-    const std::unordered_set<uintptr_t>& constant_ids_,
-    bool contiguous,
-    bool move_buffers = false);
-
 } // namespace mlx::core

mlx/backend/common/compiled_cpu.cpp

@@ -52,25 +52,8 @@ void* compile(
     return nullptr;
   }
 
-  std::string kernel_file_name;
-
-  // Deal with long kernel names. Maximum length for files on macOS is 255
-  // characters. Clip file name with a little extra room and append a 16
-  // character hash.
-  constexpr int max_file_name_length = 245;
-  if (kernel_name.size() > max_file_name_length) {
-    std::ostringstream file_name;
-    file_name
-        << std::string_view(kernel_name).substr(0, max_file_name_length - 16);
-    auto file_id = std::hash<std::string>{}(kernel_name);
-    file_name << "_" << std::hex << std::setw(16) << file_id << std::dec;
-    kernel_file_name = file_name.str();
-  } else {
-    kernel_file_name = kernel_name;
-  }
-
   std::ostringstream shared_lib_name;
-  shared_lib_name << "lib" << kernel_file_name << ".so";
+  shared_lib_name << "lib" << kernel_name << ".so";
   auto shared_lib_path = get_temp_file(shared_lib_name.str());
   bool lib_exists = false;
   {
@@ -81,7 +64,7 @@ void* compile(
   if (!lib_exists) {
     // Open source file and write source code to it
     std::ostringstream source_file_name;
-    source_file_name << kernel_file_name << ".cpp";
+    source_file_name << kernel_name << ".cpp";
     auto source_file_path = get_temp_file(source_file_name.str());
 
     std::ofstream source_file(source_file_path);
@@ -265,7 +248,28 @@ void Compiled::eval_cpu(
 
   // Figure out which kernel we are using
   auto& shape = outputs[0].shape();
-  bool contiguous = compiled_check_contiguity(inputs, shape);
+  bool contiguous = true;
+  {
+    bool all_contig = true;
+    bool all_row_contig = true;
+    bool all_col_contig = true;
+    int non_scalar_inputs = 0;
+    for (auto& x : inputs) {
+      if (is_scalar(x)) {
+        continue;
+      }
+      non_scalar_inputs++;
+      bool shape_eq = x.shape() == shape;
+      all_contig &= (x.flags().contiguous && shape_eq);
+      all_row_contig &= (x.flags().row_contiguous && shape_eq);
+      all_col_contig &= (x.flags().col_contiguous && shape_eq);
+    }
+    if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) {
+      contiguous = false;
+    } else if (non_scalar_inputs == 1 && !all_contig) {
+      contiguous = false;
+    }
+  }
 
   // Handle all broadcasting and collect function input arguments
   std::vector<void*> args;
@@ -338,8 +342,56 @@ void Compiled::eval_cpu(
     fn_ptr = compile(kernel_name, kernel.str());
   }
 
-  compiled_allocate_outputs(
-      inputs, outputs, inputs_, constant_ids_, contiguous, false);
+  // Allocate space for the outputs possibly with input donation
+  if (contiguous) {
+    int o = 0;
+    std::vector<size_t> strides;
+    size_t data_size;
+    array::Flags flags;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().contiguous && !is_scalar(in) && in.is_donatable() &&
+          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
+        outputs[o++].copy_shared_buffer(in);
+      }
+      // Get representative input flags to properly set non-donated outputs
+      if (strides.empty() && in.size() == outputs[0].size()) {
+        strides = in.strides();
+        flags = in.flags();
+        data_size = in.data_size();
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(
+          allocator::malloc_or_wait(data_size * outputs[o].itemsize()),
+          data_size,
+          strides,
+          flags);
+    }
+  } else {
+    int o = 0;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Row contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().row_contiguous && in.nbytes() == outputs[o].nbytes() &&
+          in.is_donatable() &&
+          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
+        outputs[o++].copy_shared_buffer(in);
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(allocator::malloc_or_wait(outputs[o].nbytes()));
+    }
+  }
 
   for (auto& x : outputs) {
     args.push_back(x.data<void>());

mlx/backend/common/conv.cpp

@@ -38,15 +38,11 @@ void slow_conv_1D(
 
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
   const int iH = 1 + in_dilation[0] * (in.shape(1) - 1); // Input spatial dim
-  const int C = in.shape(2); // Input channels
   const int oH = out.shape(1); // Output spatial dim
   const int O = wt.shape(0); // Out channels
+  const int C = wt.shape(2); // In channels
   const int wH = wt.shape(1); // Weight spatial dim
 
-  const int groups = C / wt.shape(2);
-  const int C_per_group = wt.shape(2);
-  const int O_per_group = O / groups;
-
   const size_t in_stride_N = in.strides()[0];
   const size_t in_stride_H = in.strides()[1];
   const size_t in_stride_C = in.strides()[2];
@@ -61,36 +57,35 @@ void slow_conv_1D(
 
   for (int n = 0; n < N; ++n) {
     for (int oh = 0; oh < oH; ++oh) {
-      for (int g = 0; g < groups; ++g) {
-        for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
-          const T* filter_wt_ptr = start_wt_ptr + o * wt_stride_O;
-          float r = 0.;
+      for (int o = 0; o < O; ++o) {
+        const T* filter_wt_ptr = start_wt_ptr + o * wt_stride_O;
+        float r = 0.;
 
-          for (int wh = 0; wh < wH; ++wh) {
-            const T* wt_ptr = filter_wt_ptr + wh * wt_stride_H;
+        for (int wh = 0; wh < wH; ++wh) {
+          const T* wt_ptr = filter_wt_ptr + wh * wt_stride_H;
 
-            int wh_flip = flip ? (wH - wh - 1) : wh;
-            int ih = oh * wt_strides[0] - padding[0] + wh_flip * wt_dilation[0];
+          int wh_flip = flip ? (wH - wh - 1) : wh;
+          int ih = oh * wt_strides[0] - padding[0] + wh_flip * wt_dilation[0];
 
-            auto ih_div = std::div(ih, in_dilation[0]);
+          auto ih_div = std::div(ih, in_dilation[0]);
 
-            if (ih >= 0 && ih < iH && ih_div.rem == 0) {
-              for (int c = g * C_per_group; c < (g + 1) * C_per_group; ++c) {
-                r += static_cast<float>(
-                         in_ptr[ih_div.quot * in_stride_H + c * in_stride_C]) *
-                    static_cast<float>(wt_ptr[(c % C_per_group) * wt_stride_C]);
-              } // c
+          if (ih >= 0 && ih < iH && ih_div.rem == 0) {
+            for (int c = 0; c < C; ++c) {
+              r += static_cast<float>(
+                       in_ptr[ih_div.quot * in_stride_H + c * in_stride_C]) *
+                  static_cast<float>(wt_ptr[c * wt_stride_C]);
+            } // c
 
-            } // ih check
-          } // wh
+          } // ih check
+        } // wh
 
-          out_ptr[oh * out_stride_H + o * out_stride_O] = static_cast<T>(r);
-        } // o
-      } // g
+        out_ptr[oh * out_stride_H + o * out_stride_O] = static_cast<T>(r);
+      } // o
     } // oh
 
     in_ptr += in_stride_N;
     out_ptr += out_stride_N;
+
   } // n
 }
 
@@ -111,17 +106,13 @@ void slow_conv_2D(
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
   const int iH = 1 + in_dilation[0] * (in.shape(1) - 1); // Input spatial dim
   const int iW = 1 + in_dilation[1] * (in.shape(2) - 1); // Input spatial dim
-  const int C = in.shape(3); // In channels
   const int oH = out.shape(1); // Output spatial dim
   const int oW = out.shape(2); // Output spatial dim
   const int O = wt.shape(0); // Out channels
+  const int C = wt.shape(3); // In channels
   const int wH = wt.shape(1); // Weight spatial dim
   const int wW = wt.shape(2); // Weight spatial dim
 
-  const int groups = C / wt.shape(3);
-  const int C_per_group = wt.shape(3);
-  const int O_per_group = O / groups;
-
   const size_t in_stride_N = in.strides()[0];
   const size_t in_stride_H = in.strides()[1];
   const size_t in_stride_W = in.strides()[2];
@@ -145,35 +136,33 @@ void slow_conv_2D(
         int ih_base = oh * wt_strides[0] - padding[0];
         int iw_base = ow * wt_strides[1] - padding[1];
 
-        for (int g = 0; g < groups; ++g) {
-          for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
-            float r = 0.;
+        for (int o = 0; o < O; ++o) {
+          float r = 0.;
 
-            for (int wh = 0; wh < wH; ++wh) {
-              for (int ww = 0; ww < wW; ++ww) {
-                int wh_flip = flip ? wH - wh - 1 : wh;
-                int ww_flip = flip ? wW - ww - 1 : ww;
-                int ih = ih_base + wh_flip * wt_dilation[0];
-                int iw = iw_base + ww_flip * wt_dilation[1];
+          for (int wh = 0; wh < wH; ++wh) {
+            for (int ww = 0; ww < wW; ++ww) {
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int ih = ih_base + wh_flip * wt_dilation[0];
+              int iw = iw_base + ww_flip * wt_dilation[1];
 
-                const T* wt_ptr_pt =
-                    wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
-                const T* in_ptr_pt =
-                    in_ptr + ih * in_stride_H + iw * in_stride_W;
+              const T* wt_ptr_pt = wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
+              const T* in_ptr_pt = in_ptr + ih * in_stride_H + iw * in_stride_W;
 
-                for (int c = g * C_per_group; c < (g + 1) * C_per_group; ++c) {
-                  r += static_cast<float>(in_ptr_pt[c * in_stride_C]) *
-                      static_cast<float>(
-                           wt_ptr_pt[(c % C_per_group) * wt_stride_C]);
-                } // c
-              } // ww
-            } // wh
+              for (int c = 0; c < C; ++c) {
+                r += static_cast<float>(in_ptr_pt[0]) *
+                    static_cast<float>(wt_ptr_pt[0]);
+                in_ptr_pt += in_stride_C;
+                wt_ptr_pt += wt_stride_C;
+              } // c
 
-            out_ptr[0] = static_cast<T>(r);
-            out_ptr += out_stride_O;
-            wt_ptr += wt_stride_O;
-          } // o
-        } // g
+            } // ww
+          } // wh
+
+          out_ptr[0] = static_cast<T>(r);
+          out_ptr += out_stride_O;
+          wt_ptr += wt_stride_O;
+        } // o
       };
 
   int jump_h = flip ? -wt_dilation[0] : wt_dilation[0];
@@ -225,43 +214,41 @@ void slow_conv_2D(
         int wh_base = base_h[oh % f_out_jump_h];
         int ww_base = base_w[ow % f_out_jump_w];
 
-        for (int g = 0; g < groups; ++g) {
-          for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
-            float r = 0.;
+        for (int o = 0; o < O; ++o) {
+          float r = 0.;
 
-            for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
-              for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
-                int wh_flip = flip ? wH - wh - 1 : wh;
-                int ww_flip = flip ? wW - ww - 1 : ww;
-                int ih = ih_base + wh_flip * wt_dilation[0];
-                int iw = iw_base + ww_flip * wt_dilation[1];
+          for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
+            for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int ih = ih_base + wh_flip * wt_dilation[0];
+              int iw = iw_base + ww_flip * wt_dilation[1];
 
-                if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) {
-                  const T* wt_ptr_pt =
-                      wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
+              if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) {
+                const T* wt_ptr_pt =
+                    wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
 
-                  int ih_dil = !is_idil_one ? (ih / in_dilation[0]) : ih;
-                  int iw_dil = !is_idil_one ? (iw / in_dilation[1]) : iw;
+                int ih_dil = !is_idil_one ? (ih / in_dilation[0]) : ih;
+                int iw_dil = !is_idil_one ? (iw / in_dilation[1]) : iw;
 
-                  const T* in_ptr_pt =
-                      in_ptr + ih_dil * in_stride_H + iw_dil * in_stride_W;
+                const T* in_ptr_pt =
+                    in_ptr + ih_dil * in_stride_H + iw_dil * in_stride_W;
 
-                  for (int c = g * C_per_group; c < (g + 1) * C_per_group;
-                       ++c) {
-                    r += static_cast<float>(in_ptr_pt[c * in_stride_C]) *
-                        static_cast<float>(
-                             wt_ptr_pt[(c % C_per_group) * wt_stride_C]);
-                  } // c
+                for (int c = 0; c < C; ++c) {
+                  r += static_cast<float>(in_ptr_pt[0]) *
+                      static_cast<float>(wt_ptr_pt[0]);
+                  in_ptr_pt += in_stride_C;
+                  wt_ptr_pt += wt_stride_C;
+                } // c
 
-                } // ih, iw check
-              } // ww
-            } // wh
+              } // ih, iw check
+            } // ww
+          } // wh
 
-            out_ptr[0] = static_cast<T>(r);
-            out_ptr += out_stride_O;
-            wt_ptr += wt_stride_O;
-          } // o
-        } // g
+          out_ptr[0] = static_cast<T>(r);
+          out_ptr += out_stride_O;
+          wt_ptr += wt_stride_O;
+        } // o
       };
 
   int oH_border_0 = 0;
@@ -318,296 +305,6 @@ void slow_conv_2D(
   } // n
 }
 
-template <typename T>
-void slow_conv_3D(
-    const array& in,
-    const array& wt,
-    array out,
-    const std::vector<int>& padding,
-    const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
-    const std::vector<int>& in_dilation,
-    bool flip) {
-  const T* st_wt_ptr = wt.data<T>();
-  const T* st_in_ptr = in.data<T>();
-  T* st_out_ptr = out.data<T>();
-
-  const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
-  const int iD = 1 + in_dilation[0] * (in.shape(1) - 1); // Input spatial dim
-  const int iH = 1 + in_dilation[1] * (in.shape(2) - 1); // Input spatial dim
-  const int iW = 1 + in_dilation[2] * (in.shape(3) - 1); // Input spatial dim
-  const int oD = out.shape(1); // Output spatial dim
-  const int oH = out.shape(2); // Output spatial dim
-  const int oW = out.shape(3); // Output spatial dim
-  const int O = wt.shape(0); // Out channels
-  const int C = wt.shape(4); // In channels
-  const int wD = wt.shape(1); // Weight spatial dim
-  const int wH = wt.shape(2); // Weight spatial dim
-  const int wW = wt.shape(3); // Weight spatial dim
-
-  const size_t in_stride_N = in.strides()[0];
-  const size_t in_stride_D = in.strides()[1];
-  const size_t in_stride_H = in.strides()[2];
-  const size_t in_stride_W = in.strides()[3];
-  const size_t in_stride_C = in.strides()[4];
-
-  const size_t wt_stride_O = wt.strides()[0];
-  const size_t wt_stride_D = wt.strides()[1];
-  const size_t wt_stride_H = wt.strides()[2];
-  const size_t wt_stride_W = wt.strides()[3];
-  const size_t wt_stride_C = wt.strides()[4];
-
-  const size_t out_stride_N = out.strides()[0];
-  const size_t out_stride_D = out.strides()[1];
-  const size_t out_stride_H = out.strides()[2];
-  const size_t out_stride_W = out.strides()[3];
-  const size_t out_stride_O = out.strides()[4];
-
-  bool is_idil_one =
-      in_dilation[0] == 1 && in_dilation[1] == 1 && in_dilation[2] == 1;
-
-  auto pt_conv_no_checks = [&](const T* in_ptr,
-                               const T* wt_ptr,
-                               T* out_ptr,
-                               int od,
-                               int oh,
-                               int ow) {
-    out_ptr += od * out_stride_D + oh * out_stride_H + ow * out_stride_W;
-    int id_base = od * wt_strides[0] - padding[0];
-    int ih_base = oh * wt_strides[1] - padding[1];
-    int iw_base = ow * wt_strides[2] - padding[2];
-
-    for (int o = 0; o < O; ++o) {
-      float r = 0.;
-
-      for (int wd = 0; wd < wD; ++wd) {
-        for (int wh = 0; wh < wH; ++wh) {
-          for (int ww = 0; ww < wW; ++ww) {
-            int wd_flip = flip ? wD - wd - 1 : wd;
-            int wh_flip = flip ? wH - wh - 1 : wh;
-            int ww_flip = flip ? wW - ww - 1 : ww;
-            int id = id_base + wd_flip * wt_dilation[0];
-            int ih = ih_base + wh_flip * wt_dilation[1];
-            int iw = iw_base + ww_flip * wt_dilation[2];
-
-            const T* wt_ptr_pt =
-                wt_ptr + wd * wt_stride_D + wh * wt_stride_H + ww * wt_stride_W;
-            const T* in_ptr_pt =
-                in_ptr + id * in_stride_D + ih * in_stride_H + iw * in_stride_W;
-
-            for (int c = 0; c < C; ++c) {
-              r += static_cast<float>(in_ptr_pt[0]) *
-                  static_cast<float>(wt_ptr_pt[0]);
-              in_ptr_pt += in_stride_C;
-              wt_ptr_pt += wt_stride_C;
-            } // c
-
-          } // ww
-        } // wh
-      } // wd
-
-      out_ptr[0] = static_cast<T>(r);
-      out_ptr += out_stride_O;
-      wt_ptr += wt_stride_O;
-    } // o
-  };
-
-  int jump_d = flip ? -wt_dilation[0] : wt_dilation[0];
-  int jump_h = flip ? -wt_dilation[1] : wt_dilation[1];
-  int jump_w = flip ? -wt_dilation[2] : wt_dilation[2];
-
-  int init_d = (flip ? (wD - 1) * wt_dilation[0] : 0);
-  int init_h = (flip ? (wH - 1) * wt_dilation[1] : 0);
-  int init_w = (flip ? (wW - 1) * wt_dilation[2] : 0);
-
-  int f_wgt_jump_d = std::lcm(in_dilation[0], wt_dilation[0]) / wt_dilation[0];
-  int f_wgt_jump_h = std::lcm(in_dilation[1], wt_dilation[1]) / wt_dilation[1];
-  int f_wgt_jump_w = std::lcm(in_dilation[2], wt_dilation[2]) / wt_dilation[2];
-
-  int f_out_jump_d = std::lcm(in_dilation[0], wt_strides[0]) / wt_strides[0];
-  int f_out_jump_h = std::lcm(in_dilation[1], wt_strides[1]) / wt_strides[1];
-  int f_out_jump_w = std::lcm(in_dilation[2], wt_strides[2]) / wt_strides[2];
-
-  std::vector<int> base_d(f_out_jump_d);
-  std::vector<int> base_h(f_out_jump_h);
-  std::vector<int> base_w(f_out_jump_w);
-
-  for (int i = 0; i < f_out_jump_d; ++i) {
-    int id_loop = i * wt_strides[0] - padding[0] + init_d;
-
-    int wd_base = 0;
-    while (wd_base < wD && id_loop % in_dilation[0] != 0) {
-      wd_base++;
-      id_loop += jump_d;
-    }
-
-    base_d[i] = wd_base;
-  }
-
-  for (int i = 0; i < f_out_jump_h; ++i) {
-    int ih_loop = i * wt_strides[1] - padding[1] + init_h;
-
-    int wh_base = 0;
-    while (wh_base < wH && ih_loop % in_dilation[1] != 0) {
-      wh_base++;
-      ih_loop += jump_h;
-    }
-
-    base_h[i] = wh_base;
-  }
-
-  for (int j = 0; j < f_out_jump_w; ++j) {
-    int iw_loop = j * wt_strides[2] - padding[2] + init_w;
-
-    int ww_base = 0;
-    while (ww_base < wW && iw_loop % in_dilation[2] != 0) {
-      ww_base++;
-      iw_loop += jump_w;
-    }
-
-    base_w[j] = ww_base;
-  }
-
-  auto pt_conv_all_checks = [&](const T* in_ptr,
-                                const T* wt_ptr,
-                                T* out_ptr,
-                                int od,
-                                int oh,
-                                int ow) {
-    out_ptr += od * out_stride_D + oh * out_stride_H + ow * out_stride_W;
-
-    int id_base = od * wt_strides[0] - padding[0];
-    int ih_base = oh * wt_strides[1] - padding[1];
-    int iw_base = ow * wt_strides[2] - padding[2];
-
-    int wd_base = base_d[od % f_out_jump_d];
-    int wh_base = base_h[oh % f_out_jump_h];
-    int ww_base = base_w[ow % f_out_jump_w];
-
-    for (int o = 0; o < O; ++o) {
-      float r = 0.;
-
-      for (int wd = wd_base; wd < wD; wd += f_wgt_jump_d) {
-        for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
-          for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
-            int wd_flip = flip ? wD - wd - 1 : wd;
-            int wh_flip = flip ? wH - wh - 1 : wh;
-            int ww_flip = flip ? wW - ww - 1 : ww;
-            int id = id_base + wd_flip * wt_dilation[0];
-            int ih = ih_base + wh_flip * wt_dilation[1];
-            int iw = iw_base + ww_flip * wt_dilation[2];
-
-            if (id >= 0 && id < iD && ih >= 0 && ih < iH && iw >= 0 &&
-                iw < iW) {
-              const T* wt_ptr_pt = wt_ptr + wd * wt_stride_D +
-                  wh * wt_stride_H + ww * wt_stride_W;
-
-              int id_dil = !is_idil_one ? (id / in_dilation[0]) : id;
-              int ih_dil = !is_idil_one ? (ih / in_dilation[1]) : ih;
-              int iw_dil = !is_idil_one ? (iw / in_dilation[2]) : iw;
-
-              const T* in_ptr_pt = in_ptr + id_dil * in_stride_D +
-                  ih_dil * in_stride_H + iw_dil * in_stride_W;
-
-              for (int c = 0; c < C; ++c) {
-                r += static_cast<float>(in_ptr_pt[0]) *
-                    static_cast<float>(wt_ptr_pt[0]);
-                in_ptr_pt += in_stride_C;
-                wt_ptr_pt += wt_stride_C;
-              } // c
-
-            } // iD, ih, iw check
-          } // ww
-        } // wh
-      } // wd
-
-      out_ptr[0] = static_cast<T>(r);
-      out_ptr += out_stride_O;
-      wt_ptr += wt_stride_O;
-    } // o
-  };
-
-  int oD_border_0 = 0;
-  int oD_border_1 =
-      is_idil_one ? ((padding[0] + wt_strides[0] - 1) / wt_strides[0]) : oD;
-  int oD_border_2 = std::max(
-      oD_border_1, (iD + padding[0] - wD * wt_dilation[0]) / wt_strides[0]);
-  int oD_border_3 = oD;
-
-  int oH_border_0 = 0;
-  int oH_border_1 =
-      is_idil_one ? ((padding[1] + wt_strides[1] - 1) / wt_strides[1]) : oH;
-  int oH_border_2 = std::max(
-      oH_border_1, (iH + padding[1] - wH * wt_dilation[1]) / wt_strides[1]);
-  int oH_border_3 = oH;
-
-  int oW_border_0 = 0;
-  int oW_border_1 =
-      is_idil_one ? ((padding[2] + wt_strides[2] - 1) / wt_strides[2]) : oW;
-  int oW_border_2 = std::max(
-      oW_border_1, (iW + padding[2] - wW * wt_dilation[2]) / wt_strides[2]);
-  int oW_border_3 = oW;
-
-  for (int n = 0; n < N; ++n) {
-    // Case 1: od might put us out of bounds
-    for (int od = oD_border_0; od < oD_border_1; ++od) {
-      for (int oh = 0; oh < oH; ++oh) {
-        for (int ow = 0; ow < oW; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-      } // oh
-    } // od
-
-    // Case 2: od in bounds
-    for (int od = oD_border_1; od < oD_border_2; ++od) {
-      // Case 2.1: oh might put us out of bounds
-      for (int oh = oH_border_0; oh < oH_border_1; ++oh) {
-        for (int ow = 0; ow < oW; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-      } // oh
-
-      // Case 2.2: oh in bounds
-      for (int oh = oH_border_1; oh < oH_border_2; ++oh) {
-        // Case 2.2.1: ow might put us out of bounds
-        for (int ow = oW_border_0; ow < oW_border_1; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-
-        // Case 2.2.2: ow in bounds
-        for (int ow = oW_border_1; ow < oW_border_2; ++ow) {
-          pt_conv_no_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-
-        // Case 2.2.3: ow might put us out of bounds
-        for (int ow = oW_border_2; ow < oW_border_3; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-      } // oh
-
-      // Case 2.3: oh might put us out of bounds
-      for (int oh = oH_border_2; oh < oH_border_3; ++oh) {
-        for (int ow = 0; ow < oW; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-      } // oh
-    } // od
-
-    // Case 3: od might put us out of bounds
-    for (int od = oD_border_2; od < oD_border_3; ++od) {
-      for (int oh = 0; oh < oH; ++oh) {
-        for (int ow = 0; ow < oW; ++ow) {
-          pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
-        } // ow
-      } // oh
-    } // od
-
-    st_in_ptr += in_stride_N;
-    st_out_ptr += out_stride_N;
-
-  } // n
-}
-
 void dispatch_slow_conv_1D(
     const array& in,
     const array& wt,
@@ -656,30 +353,6 @@ void dispatch_slow_conv_2D(
   }
 }
 
-void dispatch_slow_conv_3D(
-    const array& in,
-    const array& wt,
-    array out,
-    const std::vector<int>& padding,
-    const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
-    const std::vector<int>& in_dilation,
-    bool flip) {
-  if (in.dtype() == float32) {
-    return slow_conv_3D<float>(
-        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
-  } else if (in.dtype() == float16) {
-    return slow_conv_3D<float16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
-  } else if (in.dtype() == bfloat16) {
-    return slow_conv_3D<bfloat16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
-  } else {
-    throw std::invalid_argument(
-        "[Convolution::eval] got unsupported data type.");
-  }
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Explicit gemm conv
 ///////////////////////////////////////////////////////////////////////////////
@@ -693,15 +366,11 @@ void explicit_gemm_conv_1D_cpu(
     const std::vector<int>& wt_dilation) {
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
   const int iH = in.shape(1); // Input spatial dim
-  const int C = in.shape(2); // Input channels
   const int oH = out.shape(1); // Output spatial dim
   const int O = wt.shape(0); // Out channels
+  const int C = wt.shape(2); // In channels
   const int wH = wt.shape(1); // Weight spatial dim
 
-  const int groups = C / wt.shape(2);
-  const int C_per_group = wt.shape(2);
-  const int O_per_group = O / groups;
-
   auto conv_dtype = float32;
 
   // Pad input
@@ -733,11 +402,6 @@ void explicit_gemm_conv_1D_cpu(
       in_padded.strides()[1],
       in_padded.strides()[2]};
   auto flags = in_padded.flags();
-  if (groups > 1) {
-    // Transpose the last two dimensions for grouped convolutions
-    std::swap(strided_shape[2], strided_shape[3]);
-    std::swap(strided_strides[2], strided_strides[3]);
-  }
 
   array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
   in_strided_view.copy_shared_buffer(
@@ -752,19 +416,7 @@ void explicit_gemm_conv_1D_cpu(
   auto gemm_wt = wt;
   auto gemm_out = out;
 
-  if (groups > 1) {
-    // Transpose the last two dimensions for grouped convolutions
-    array wt_transpose(
-        {wt.shape(0), wt.shape(2), wt.shape(1)}, wt.dtype(), nullptr, {});
-    wt_transpose.copy_shared_buffer(
-        wt,
-        {wt.strides(0), wt.strides(2), wt.strides(1)},
-        wt.flags(),
-        wt.size(),
-        0);
-    gemm_wt = array(wt_transpose.shape(), float32, nullptr, {});
-    copy(wt_transpose, gemm_wt, CopyType::General);
-  } else if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
+  if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
     auto ctype =
         wt.flags().row_contiguous ? CopyType::Vector : CopyType::General;
     gemm_wt = array(wt.shape(), float32, nullptr, {});
@@ -776,29 +428,27 @@ void explicit_gemm_conv_1D_cpu(
     gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
   }
 
-  for (int g = 0; g < groups; ++g) {
-    // Perform gemm
-    cblas_sgemm(
-        CblasRowMajor,
-        CblasNoTrans, // no trans A
-        CblasTrans, // transB
-        strided_reshape[0], // M
-        O_per_group, // N
-        C_per_group * wH, // K
-        1.0f, // alpha
-        in_strided.data<float>() + g * C_per_group * wH, // A
-        wH * C, // lda
-        gemm_wt.data<float>() + g * O_per_group * C_per_group * wH, // B
-        wH * C_per_group, // ldb
-        0.0f, // beta
-        gemm_out.data<float>() + g * O_per_group, // C
-        O // ldc
-    );
+  // Perform gemm
+  cblas_sgemm(
+      CblasRowMajor,
+      CblasNoTrans, // no trans A
+      CblasTrans, // transB
+      strided_reshape[0], // M
+      O, // N
+      strided_reshape[1], // K
+      1.0f, // alpha
+      in_strided.data<float>(),
+      strided_reshape[1], // lda
+      gemm_wt.data<float>(),
+      strided_reshape[1], // ldb
+      0.0f, // beta
+      gemm_out.data<float>(),
+      O // ldc
+  );
 
-    // Copy results if needed
-    if (out.dtype() != float32) {
-      copy(gemm_out, out, CopyType::Vector);
-    }
+  // Copy results if needed
+  if (out.dtype() != float32) {
+    copy(gemm_out, out, CopyType::Vector);
   }
 }
 
@@ -904,131 +554,6 @@ void explicit_gemm_conv_2D_cpu(
   }
 }
 
-void explicit_gemm_conv_ND_cpu(
-    const array& in,
-    const array& wt,
-    array out,
-    const std::vector<int>& padding,
-    const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
-  const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
-  const auto iDim = std::vector<int>(
-      in.shape().begin() + 1, in.shape().end() - 1); // Input spatial dim
-  const auto oDim = std::vector<int>(
-      out.shape().begin() + 1, out.shape().end() - 1); // Output spatial dim
-  const int O = wt.shape(0); // Out channels
-  const int C = wt.shape(-1); // In channels
-  const auto wDim = std::vector<int>(
-      wt.shape().begin() + 1, wt.shape().end() - 1); // Weight spatial dim
-
-  auto conv_dtype = float32;
-
-  // Pad input
-  std::vector<int> padded_shape(in.shape().size());
-  padded_shape.front() = N;
-  for (size_t i = 0; i < iDim.size(); i++) {
-    padded_shape[i + 1] = iDim[i] + 2 * padding[i];
-  }
-  padded_shape.back() = C;
-  array in_padded(padded_shape, conv_dtype, nullptr, {});
-
-  // Fill with zeros
-  copy(array(0, conv_dtype), in_padded, CopyType::Scalar);
-
-  // Pick input slice from padded
-  size_t data_offset = 0;
-  for (size_t i = 0; i < padding.size(); i++) {
-    data_offset += padding[i] * in_padded.strides()[i + 1];
-  }
-  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
-  in_padded_slice.copy_shared_buffer(
-      in_padded,
-      in_padded.strides(),
-      in_padded.flags(),
-      in_padded_slice.size(),
-      data_offset);
-
-  // Copy input values into the slice
-  copy_inplace(in, in_padded_slice, CopyType::GeneralGeneral);
-
-  // Make strided view
-  std::vector<int> strided_shape(oDim.size() + wDim.size() + 2);
-  strided_shape.front() = N;
-  for (size_t i = 0; i < oDim.size(); i++) {
-    strided_shape[i + 1] = oDim[i];
-  }
-  for (size_t i = 0; i < wDim.size(); i++) {
-    strided_shape[i + 1 + oDim.size()] = wDim[i];
-  }
-  strided_shape.back() = C;
-
-  std::vector<size_t> strided_strides(in.shape().size() * 2 - 2);
-  strided_strides[0] = in_padded.strides()[0];
-  for (size_t i = 0; i < wt_strides.size(); i++) {
-    strided_strides[i + 1] = in_padded.strides()[i + 1] * wt_strides[i];
-  }
-  for (size_t i = 1; i < in_padded.strides().size(); i++) {
-    strided_strides[i + wt_strides.size()] = in_padded.strides()[i];
-  }
-
-  auto flags = in_padded.flags();
-
-  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
-  in_strided_view.copy_shared_buffer(
-      in_padded, strided_strides, flags, in_strided_view.size(), 0);
-
-  // Materialize strided view
-  std::vector<int> strided_reshape = {N, C};
-  for (const auto& o : oDim) {
-    strided_reshape[0] *= o;
-  }
-  for (const auto& w : wDim) {
-    strided_reshape[1] *= w;
-  }
-
-  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
-  copy(in_strided_view, in_strided, CopyType::General);
-
-  // Check wt dtype and prepare
-  auto gemm_wt = wt;
-  auto gemm_out = out;
-
-  if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
-    auto ctype =
-        wt.flags().row_contiguous ? CopyType::Vector : CopyType::General;
-    gemm_wt = array(wt.shape(), float32, nullptr, {});
-    copy(wt, gemm_wt, ctype);
-  }
-
-  if (out.dtype() != float32) {
-    gemm_out = array(out.shape(), float32, nullptr, {});
-    gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
-  }
-
-  // Perform gemm
-  cblas_sgemm(
-      CblasRowMajor,
-      CblasNoTrans, // no trans A
-      CblasTrans, // transB
-      strided_reshape[0], // M
-      O, // N
-      strided_reshape[1], // K
-      1.0f, // alpha
-      in_strided.data<float>(),
-      strided_reshape[1], // lda
-      gemm_wt.data<float>(),
-      strided_reshape[1], // ldb
-      0.0f, // beta
-      gemm_out.data<float>(),
-      O // ldc
-  );
-
-  // Copy results if needed
-  if (out.dtype() != float32) {
-    copy(gemm_out, out, CopyType::Vector);
-  }
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Conv routing
 ///////////////////////////////////////////////////////////////////////////////
@@ -1064,19 +589,6 @@ void conv_2D_cpu(
       in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
 }
 
-void conv_3D_cpu(
-    const array& in,
-    const array& wt,
-    array out,
-    const std::vector<int>& padding,
-    const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation,
-    const std::vector<int>& in_dilation,
-    bool flip) {
-  return dispatch_slow_conv_3D(
-      in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
-}
-
 } // namespace
 
 void Convolution::eval(const std::vector<array>& inputs, array& out) {
@@ -1085,20 +597,8 @@ void Convolution::eval(const std::vector<array>& inputs, array& out) {
   auto& in = inputs[0];
   auto& wt = inputs[1];
 
-  // 3D convolution
-  if (in.ndim() == (3 + 2)) {
-    return conv_3D_cpu(
-        in,
-        wt,
-        out,
-        padding_,
-        kernel_strides_,
-        kernel_dilation_,
-        input_dilation_,
-        flip_);
-  }
   // 2D convolution
-  else if (in.ndim() == (2 + 2)) {
+  if (in.ndim() == (2 + 2)) {
     return conv_2D_cpu(
         in,
         wt,

mlx/backend/common/copy.cpp

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

mlx/backend/common/copy.h

@@ -1,4 +1,4 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #pragma once
 
@@ -26,15 +26,4 @@ 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

@@ -5,6 +5,7 @@
 #else
 #include <cblas.h>
 #endif
+
 #include <cstring>
 
 #include "mlx/array.h"
@@ -33,7 +34,6 @@ DEFAULT(ArcCosh)
 DEFAULT(ArcSin)
 DEFAULT(ArcSinh)
 DEFAULT(ArcTan)
-DEFAULT(ArcTan2)
 DEFAULT(ArcTanh)
 DEFAULT(ArgPartition)
 DEFAULT(ArgReduce)
@@ -41,34 +41,27 @@ DEFAULT(ArgSort)
 DEFAULT(AsType)
 DEFAULT(AsStrided)
 DEFAULT(Broadcast)
-DEFAULT(BlockMaskedMM)
-DEFAULT(GatherMM)
-DEFAULT(GatherQMM)
 DEFAULT_MULTI(DivMod)
 DEFAULT(Ceil)
 DEFAULT(Concatenate)
-DEFAULT(Conjugate)
 DEFAULT(Convolution)
 DEFAULT(Copy)
 DEFAULT(Cos)
 DEFAULT(Cosh)
-DEFAULT_MULTI(CustomTransforms)
+DEFAULT_MULTI(CustomVJP)
 DEFAULT_MULTI(Depends)
 DEFAULT(Divide)
-DEFAULT(NumberOfElements)
 DEFAULT(Remainder)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
 DEFAULT(Exp)
-DEFAULT(Expm1)
 DEFAULT(FFT)
 DEFAULT(Floor)
 DEFAULT(Full)
 DEFAULT(Gather)
 DEFAULT(Greater)
 DEFAULT(GreaterEqual)
-DEFAULT(Hadamard)
 DEFAULT(Less)
 DEFAULT(LessEqual)
 DEFAULT(Load)
@@ -100,7 +93,6 @@ DEFAULT(Sign)
 DEFAULT(Sin)
 DEFAULT(Sinh)
 DEFAULT(Slice)
-DEFAULT(SliceUpdate)
 DEFAULT(Softmax)
 DEFAULT(Sort)
 DEFAULT_MULTI(Split)
@@ -108,12 +100,9 @@ DEFAULT(Square)
 DEFAULT(Sqrt)
 DEFAULT(StopGradient)
 DEFAULT(Subtract)
-DEFAULT_MULTI(SVD)
 DEFAULT(Tan)
 DEFAULT(Tanh)
 DEFAULT(Transpose)
-DEFAULT(Inverse)
-DEFAULT(Cholesky)
 
 namespace {
 

mlx/backend/common/hadamard.cpp

@@ -1,107 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include <cassert>
-
-#include "mlx/backend/common/copy.h"
-#include "mlx/backend/common/hadamard.h"
-#include "mlx/primitives.h"
-
-namespace mlx::core {
-
-// n = 2^k component
-template <typename T>
-void hadamard_n(array& out, int n, int m, float scale) {
-  for (int b = 0; b < out.size() / n; b++) {
-    size_t loc = b * n;
-    T* data_ptr = out.data<T>() + loc;
-    int h = 1;
-    int n_over_2 = n / 2;
-    while (h < n) {
-      for (int i = 0; i < n / 2; i++) {
-        int k = i & (h - 1);
-        int j = ((i - k) << 1) + k;
-        float x = *(data_ptr + j);
-        float y = *(data_ptr + j + h);
-        *(data_ptr + j) = x + y;
-        *(data_ptr + j + h) = x - y;
-        if (h == n_over_2) {
-          *(data_ptr + j) *= scale;
-          *(data_ptr + j + h) *= scale;
-        }
-      }
-      h <<= 1;
-    }
-  }
-}
-
-// m component
-template <typename T>
-void hadamard_m(array& out, int n, int m, float scale) {
-  auto h_matrices = hadamard_matrices();
-  auto& matrix = h_matrices[m];
-  auto start = 1;
-  auto end = matrix.find('\n', start);
-  std::vector<bool> hmat_vec;
-  while (end != std::string_view::npos) {
-    auto row = matrix.substr(start, end - start);
-    for (int i = 0; i < row.length(); i++) {
-      hmat_vec.push_back(row[i] == '+');
-    }
-    start = end + 1;
-    end = matrix.find('\n', start);
-  }
-
-  for (int b = 0; b < out.size() / m / n; b++) {
-    size_t loc = b * n * m;
-    T* data_ptr = out.data<T>() + loc;
-    for (int i = 0; i < n; i++) {
-      std::vector<float> out(m);
-      for (int j = 0; j < m; j++) {
-        for (int k = 0; k < m; k++) {
-          float x = *(data_ptr + i + k * n);
-          if (hmat_vec[k + j * m]) {
-            out[j] += x;
-          } else {
-            out[j] -= x;
-          }
-        }
-      }
-      for (int j = 0; j < m; j++) {
-        *(data_ptr + i + j * n) = out[j] * scale;
-      }
-    }
-  }
-}
-
-template <typename T>
-void hadamard(array& out, int n, int m, float scale) {
-  float n_scale = m > 1 ? 1.0 : scale;
-  hadamard_n<T>(out, n, m, n_scale);
-  if (m > 1) {
-    hadamard_m<T>(out, n, m, scale);
-  }
-}
-
-void Hadamard::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  auto& in = inputs[0];
-
-  // Copy input to output
-  copy(in, out, CopyType::General);
-
-  int axis = out.ndim() - 1;
-  auto [n, m] = decompose_hadamard(out.shape(axis));
-
-  switch (in.dtype()) {
-    case float32:
-      return hadamard<float>(out, n, m, scale_);
-    case float16:
-      return hadamard<float16_t>(out, n, m, scale_);
-    case bfloat16:
-      return hadamard<bfloat16_t>(out, n, m, scale_);
-    default:
-      throw std::invalid_argument("[hadamard] Unsupported type.");
-  }
-}
-
-} // namespace mlx::core

mlx/backend/common/hadamard.h

@@ -1,105 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#include <map>
-
-#include "mlx/utils.h"
-
-namespace mlx::core {
-
-// From http://neilsloane.com/hadamard/
-constexpr std::string_view h12 = R"(
-+-++++++++++
---+-+-+-+-+-
-+++-++----++
-+---+--+-++-
-+++++-++----
-+-+---+--+-+
-++--+++-++--
-+--++---+--+
-++----+++-++
-+--+-++---+-
-++++----+++-
-+-+--+-++---
-)";
-
-constexpr std::string_view h20 = R"(
-+----+----++--++-++-
--+----+---+++---+-++
---+----+---+++-+-+-+
----+----+---+++++-+-
-----+----++--++-++-+
--+++++-----+--+++--+
-+-+++-+---+-+--+++--
-++-++--+---+-+--+++-
-+++-+---+---+-+--+++
-++++-----++--+-+--++
---++-+-++-+-----++++
----++-+-++-+---+-+++
-+---++-+-+--+--++-++
-++---++-+----+-+++-+
--++---++-+----+++++-
--+--+--++-+----+----
-+-+-----++-+----+---
--+-+-+---+--+----+--
---+-+++------+----+-
-+--+--++------+----+
-)";
-
-constexpr std::string_view h28 = R"(
-+------++----++-+--+-+--++--
--+-----+++-----+-+--+-+--++-
---+-----+++---+-+-+----+--++
----+-----+++---+-+-+-+--+--+
-----+-----+++---+-+-+++--+--
------+-----++++--+-+--++--+-
-------++----++-+--+-+--++--+
---++++-+-------++--+++-+--+-
----++++-+-----+-++--+-+-+--+
-+---+++--+----++-++--+-+-+--
-++---++---+----++-++--+-+-+-
-+++---+----+----++-++--+-+-+
-++++--------+-+--++-++--+-+-
--++++--------+++--++--+--+-+
--+-++-++--++--+--------++++-
-+-+-++--+--++--+--------++++
--+-+-++--+--++--+----+---+++
-+-+-+-++--+--+---+---++---++
-++-+-+-++--+------+--+++---+
--++-+-+-++--+------+-++++---
-+-++-+---++--+------+-++++--
--++--++-+-++-+++----++------
-+-++--++-+-++-+++-----+-----
-++-++---+-+-++-+++-----+----
--++-++-+-+-+-+--+++-----+---
---++-++++-+-+----+++-----+--
-+--++-+-++-+-+----+++-----+-
-++--++-+-++-+-+----++------+
-)";
-
-inline const std::map<int, std::string_view> hadamard_matrices() {
-  return {{12, h12}, {20, h20}, {28, h28}};
-}
-
-inline std::pair<int, int> decompose_hadamard(int n) {
-  // n = m*2^k
-  int m = 1;
-  if (!is_power_of_2(n)) {
-    auto h_matrices = hadamard_matrices();
-    for (auto [factor, _] : h_matrices) {
-      if (n % factor == 0) {
-        m = factor;
-        n /= factor;
-        break;
-      }
-    }
-    if (m == 1) {
-      throw std::invalid_argument(
-          "[hadamard] Only supports n = m*2^k where m in (1, 12, 20, 28).");
-    }
-  }
-  return {n, m};
-}
-
-} // namespace mlx::core

mlx/backend/common/inverse.cpp

@@ -1,140 +0,0 @@
-// Copyright © 2023-2024 Apple Inc.
-
-#include "mlx/allocator.h"
-#include "mlx/backend/common/copy.h"
-#include "mlx/primitives.h"
-
-#ifdef ACCELERATE_NEW_LAPACK
-#include <Accelerate/Accelerate.h>
-#else
-#include <lapack.h>
-#endif
-
-// Wrapper to account for differences in
-// LAPACK implementations (basically how to pass the 'uplo' string to fortran).
-int strtri_wrapper(char uplo, char diag, float* matrix, int N) {
-  int info;
-
-#ifdef LAPACK_FORTRAN_STRLEN_END
-  strtri_(
-      /* uplo = */ &uplo,
-      /* diag = */ &diag,
-      /* N = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info,
-      /* uplo_len = */ static_cast<size_t>(1),
-      /* diag_len = */ static_cast<size_t>(1));
-#else
-  strtri_(
-      /* uplo = */ &uplo,
-      /* diag = */ &diag,
-      /* N = */ &N,
-      /* a = */ matrix,
-      /* lda = */ &N,
-      /* info = */ &info);
-#endif
-
-  return info;
-}
-
-namespace mlx::core {
-
-void general_inv(array& inv, int N, int i) {
-  int info;
-  auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
-  // 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 tri_inv(array& inv, int N, int i, bool upper) {
-  const char uplo = upper ? 'L' : 'U';
-  const char diag = 'N';
-  int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
-  if (info != 0) {
-    std::stringstream ss;
-    ss << "inverse_impl: triangular inversion failed with error code " << info;
-    throw std::runtime_error(ss.str());
-  }
-}
-
-void inverse_impl(const array& a, array& inv, bool tri, bool upper) {
-  // 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);
-
-  for (int i = 0; i < num_matrices; i++) {
-    if (tri) {
-      tri_inv(inv, N, i, upper);
-    } else {
-      general_inv(inv, N, i);
-    }
-  }
-}
-
-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, tri_, upper_);
-}
-
-} // namespace mlx::core

mlx/backend/common/lapack_helper.h

@@ -1,23 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#ifdef ACCELERATE_NEW_LAPACK
-#include <Accelerate/Accelerate.h>
-#else
-#include <lapack.h>
-#endif
-
-#if defined(LAPACK_GLOBAL) || defined(LAPACK_NAME)
-
-// This is to work around a change in the function signatures of lapack >= 3.9.1
-// where functions taking char* also include a strlen argument, see a similar
-// change in OpenCV:
-// https://github.com/opencv/opencv/blob/1eb061f89de0fb85c4c75a2deeb0f61a961a63ad/cmake/OpenCVFindLAPACK.cmake#L57
-#define MLX_LAPACK_FUNC(f) LAPACK_##f
-
-#else
-
-#define MLX_LAPACK_FUNC(f) f##_
-
-#endif

mlx/backend/common/make_compiled_preamble.sh

@@ -11,7 +11,7 @@ GCC=$2
 SRCDIR=$3
 CLANG=$4
 
-if [ "$CLANG" = "TRUE" ]; then
+if [ $CLANG = "TRUE" ]; then
   read -r -d '' INCLUDES <<- EOM
   #include <cmath>
   #include <complex>
@@ -28,7 +28,6 @@ const char* get_kernel_preamble() {
 return R"preamble(
 $INCLUDES
 $CONTENT
-using namespace mlx::core;
 using namespace mlx::core::detail;
 )preamble";
 }

mlx/backend/common/masked_mm.cpp

@@ -1,305 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#ifdef ACCELERATE_NEW_LAPACK
-#include <Accelerate/Accelerate.h>
-#else
-#include <cblas.h>
-#endif
-
-#include <cstring>
-
-#include "mlx/array.h"
-#include "mlx/backend/common/copy.h"
-#include "mlx/backend/common/utils.h"
-#include "mlx/primitives.h"
-
-namespace mlx::core {
-
-namespace {
-
-template <typename T, typename mask_t>
-inline void mask_matrix(
-    T* data,
-    const mask_t* mask,
-    int block_size,
-    const int X,
-    const int Y,
-    const size_t X_data_str,
-    const size_t Y_data_str,
-    const size_t X_mask_str,
-    const size_t Y_mask_str,
-    const size_t mask_offset) {
-  int tX = (X + block_size - 1) / block_size;
-  int tY = (Y + block_size - 1) / block_size;
-
-  for (int i = 0; i < tX; i++) {
-    for (int j = 0; j < tY; j++) {
-      mask_t do_mask = mask[mask_offset + i * X_mask_str + j * Y_mask_str];
-      if (do_mask != 1) {
-        int loc_x = i * block_size;
-        int loc_y = j * block_size;
-        T* data_block = data + loc_x * X_data_str + loc_y * Y_data_str;
-
-        int size_x = std::min(block_size, X - loc_x);
-        int size_y = std::min(block_size, Y - loc_y);
-        for (int ii = 0; ii < size_x; ii++) {
-          for (int jj = 0; jj < size_y; jj++) {
-            if constexpr (std::is_same_v<mask_t, bool>) {
-              data_block[ii * X_data_str + jj * Y_data_str] = T(0.);
-            } else {
-              data_block[ii * X_data_str + jj * Y_data_str] *= do_mask;
-            }
-          }
-        }
-      }
-    }
-  }
-}
-
-} // namespace
-
-void BlockMaskedMM::eval(const std::vector<array>& inputs, array& out) {
-  if (out.dtype() != float32) {
-    throw std::runtime_error(
-        "[BlockMaskedMM::eval] Currently only supports float32.");
-  }
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  auto& a_pre = inputs[0];
-  auto& b_pre = inputs[1];
-
-  auto check_transpose =
-      [](const array& arr, bool do_copy, bool expand_all = false) {
-        auto stx = arr.strides()[arr.ndim() - 2];
-        auto sty = arr.strides()[arr.ndim() - 1];
-        if (!expand_all && stx == arr.shape(-1) && sty == 1) {
-          if (do_copy) {
-            array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-            copy(arr, arr_copy, CopyType::Vector);
-            return std::make_tuple(false, stx, arr_copy);
-          }
-          return std::make_tuple(false, stx, arr);
-        } else if (!expand_all && stx == 1 && sty == arr.shape(-2)) {
-          if (do_copy) {
-            array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-            copy(arr, arr_copy, CopyType::Vector);
-            return std::make_tuple(true, sty, arr_copy);
-          }
-          return std::make_tuple(true, sty, arr);
-        } else {
-          array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-          copy(arr, arr_copy, CopyType::General);
-          size_t stx = arr.shape(-1);
-          return std::make_tuple(false, stx, arr_copy);
-        }
-      };
-
-  bool has_op_mask = inputs.size() > 3;
-  bool has_out_mask = inputs.size() == 3 || inputs.size() == 5;
-  auto [a_transposed, lda, a] =
-      check_transpose(a_pre, has_op_mask, inputs.back().dtype() != bool_);
-  auto [b_transposed, ldb, b] =
-      check_transpose(b_pre, has_op_mask, inputs.back().dtype() != bool_);
-
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
-
-  if (M == 0 || N == 0) {
-    return;
-  }
-
-  if (K == 0) {
-    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
-    return;
-  }
-
-  auto mask_array = [](const array& mask,
-                       float* data,
-                       int block_size,
-                       int batch_idx,
-                       int X,
-                       int Y,
-                       size_t X_data_str,
-                       size_t Y_data_str) {
-    size_t mask_offset = elem_to_loc(
-        mask.shape(-1) * mask.shape(-2) * batch_idx,
-        mask.shape(),
-        mask.strides());
-
-    size_t X_mask_str = mask.strides()[mask.ndim() - 2];
-    size_t Y_mask_str = mask.strides()[mask.ndim() - 1];
-
-    if (mask.dtype() == bool_) {
-      return mask_matrix(
-          data,
-          mask.data<bool>(),
-          block_size,
-          X,
-          Y,
-          X_data_str,
-          Y_data_str,
-          X_mask_str,
-          Y_mask_str,
-          mask_offset);
-    } else {
-      return mask_matrix(
-          data,
-          mask.data<float>(),
-          block_size,
-          X,
-          Y,
-          X_data_str,
-          Y_data_str,
-          X_mask_str,
-          Y_mask_str,
-          mask_offset);
-    }
-  };
-
-  for (int i = 0; i < (out.size() / (M * size_t(N))); ++i) {
-    // Adjust pointer
-    float* ai =
-        a.data<float>() + elem_to_loc(M * K * i, a.shape(), a.strides());
-    float* bi =
-        b.data<float>() + elem_to_loc(K * N * i, b.shape(), b.strides());
-    float* ci = out.data<float>() + M * N * i;
-
-    // Zero out blocks in a and b if needed
-    if (has_op_mask) {
-      auto& a_mask = inputs[inputs.size() - 2];
-      mask_array(
-          a_mask,
-          ai,
-          block_size_,
-          i,
-          M,
-          K,
-          a_transposed ? 1 : lda,
-          a_transposed ? lda : 1);
-
-      auto& b_mask = inputs[inputs.size() - 1];
-      mask_array(
-          b_mask,
-          bi,
-          block_size_,
-          i,
-          K,
-          N,
-          b_transposed ? 1 : ldb,
-          b_transposed ? ldb : 1);
-    }
-
-    // Do matmul
-    cblas_sgemm(
-        CblasRowMajor,
-        a_transposed ? CblasTrans : CblasNoTrans, // transA
-        b_transposed ? CblasTrans : CblasNoTrans, // transB
-        M,
-        N,
-        K,
-        1.0, // alpha
-        ai,
-        lda,
-        bi,
-        ldb,
-        0.0, // beta
-        ci,
-        out.shape(-1) // ldc
-    );
-
-    // Zero out blocks in out
-    if (has_out_mask) {
-      mask_array(inputs[2], ci, block_size_, i, M, N, N, 1);
-    }
-  }
-}
-
-void GatherMM::eval(const std::vector<array>& inputs, array& out) {
-  if (out.dtype() != float32) {
-    throw std::runtime_error(
-        "[GatherMM::eval] Currently only supports float32.");
-  }
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-
-  auto& a_pre = inputs[0];
-  auto& b_pre = inputs[1];
-
-  auto check_transpose = [](const array& arr) {
-    auto stx = arr.strides()[arr.ndim() - 2];
-    auto sty = arr.strides()[arr.ndim() - 1];
-    if (stx == arr.shape(-1) && sty == 1) {
-      return std::make_tuple(false, stx, arr);
-    } else if (stx == 1 && sty == arr.shape(-2)) {
-      return std::make_tuple(true, sty, arr);
-    } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
-      size_t stx = arr.shape(-1);
-      return std::make_tuple(false, stx, arr_copy);
-    }
-  };
-
-  auto [a_transposed, lda, a] = check_transpose(a_pre);
-  auto [b_transposed, ldb, b] = check_transpose(b_pre);
-
-  size_t M = a.shape(-2);
-  size_t N = b.shape(-1);
-  size_t K = a.shape(-1);
-
-  if (M == 0 || N == 0) {
-    return;
-  }
-
-  if (K == 0) {
-    std::memset(static_cast<void*>(out.data<float>()), 0, out.nbytes());
-    return;
-  }
-
-  // Get batch dims
-  auto batch_size_out = out.size() / (M * N);
-  size_t matrix_stride_out = M * N;
-
-  auto get_batch_dims = [](const auto& v) {
-    return decltype(v){v.begin(), v.end() - 2};
-  };
-
-  auto& lhs_indices = inputs[2];
-  auto& rhs_indices = inputs[3];
-
-  std::vector<int> batch_shape = get_batch_dims(out.shape());
-  int batch_ndim = batch_shape.size();
-
-  std::vector<int> batch_shape_A = get_batch_dims(a.shape());
-  std::vector<size_t> batch_strides_A = get_batch_dims(a.strides());
-  std::vector<int> batch_shape_B = get_batch_dims(b.shape());
-  std::vector<size_t> batch_strides_B = get_batch_dims(b.strides());
-
-  const uint32_t* lhs_indices_ptr = lhs_indices.data<uint32_t>();
-  const uint32_t* rhs_indices_ptr = rhs_indices.data<uint32_t>();
-
-  for (int i = 0; i < batch_size_out; i++) {
-    // Get index
-    uint32_t indx_A = lhs_indices_ptr[elem_to_loc(i, lhs_indices)];
-    uint32_t indx_B = rhs_indices_ptr[elem_to_loc(i, rhs_indices)];
-
-    cblas_sgemm(
-        CblasRowMajor,
-        a_transposed ? CblasTrans : CblasNoTrans, // transA
-        b_transposed ? CblasTrans : CblasNoTrans, // transB
-        M,
-        N,
-        K,
-        1.0f, // alpha
-        a.data<float>() + elem_to_loc(indx_A, batch_shape_A, batch_strides_A),
-        lda,
-        b.data<float>() + elem_to_loc(indx_B, batch_shape_B, batch_strides_B),
-        ldb,
-        0.0f, // beta
-        out.data<float>() + matrix_stride_out * i,
-        out.shape(-1) // ldc
-    );
-  }
-}
-
-} // namespace mlx::core

mlx/backend/common/ops.h

@@ -108,233 +108,213 @@ struct Abs {
   template <typename T>
   T operator()(T x) {
     return std::abs(x);
-  }
+  };
   uint8_t operator()(uint8_t x) {
     return x;
-  }
+  };
   uint16_t operator()(uint16_t x) {
     return x;
-  }
+  };
   uint32_t operator()(uint32_t x) {
     return x;
-  }
+  };
   uint64_t operator()(uint64_t x) {
     return x;
-  }
+  };
   bool operator()(bool x) {
     return x;
-  }
+  };
 };
 
 struct ArcCos {
   template <typename T>
   T operator()(T x) {
     return std::acos(x);
-  }
+  };
 };
 
 struct ArcCosh {
   template <typename T>
   T operator()(T x) {
     return std::acosh(x);
-  }
+  };
 };
 
 struct ArcSin {
   template <typename T>
   T operator()(T x) {
     return std::asin(x);
-  }
+  };
 };
 
 struct ArcSinh {
   template <typename T>
   T operator()(T x) {
     return std::asinh(x);
-  }
+  };
 };
 
 struct ArcTan {
   template <typename T>
   T operator()(T x) {
     return std::atan(x);
-  }
-};
-
-struct ArcTan2 {
-  template <typename T>
-  T operator()(T y, T x) {
-    return std::atan2(y, x);
-  }
+  };
 };
 
 struct ArcTanh {
   template <typename T>
   T operator()(T x) {
     return std::atanh(x);
-  }
+  };
 };
 
 struct Ceil {
   template <typename T>
   T operator()(T x) {
     return std::ceil(x);
-  }
+  };
   int8_t operator()(int8_t x) {
     return x;
-  }
+  };
   int16_t operator()(int16_t x) {
     return x;
-  }
+  };
   int32_t operator()(int32_t x) {
     return x;
-  }
+  };
   int64_t operator()(int64_t x) {
     return x;
-  }
+  };
   uint8_t operator()(uint8_t x) {
     return x;
-  }
+  };
   uint16_t operator()(uint16_t x) {
     return x;
-  }
+  };
   uint32_t operator()(uint32_t x) {
     return x;
-  }
+  };
   uint64_t operator()(uint64_t x) {
     return x;
-  }
+  };
   bool operator()(bool x) {
     return x;
-  }
-};
-
-struct Conjugate {
-  complex64_t operator()(complex64_t x) {
-    return std::conj(x);
-  }
+  };
 };
 
 struct Cos {
   template <typename T>
   T operator()(T x) {
     return std::cos(x);
-  }
+  };
 };
 
 struct Cosh {
   template <typename T>
   T operator()(T x) {
     return std::cosh(x);
-  }
+  };
 };
 
 struct Erf {
   template <typename T>
   T operator()(T x) {
     return static_cast<T>(fast_erf(static_cast<float>(x)));
-  }
+  };
 };
 
 struct ErfInv {
   template <typename T>
   T operator()(T x) {
     return static_cast<T>(fast_erfinv(static_cast<float>(x)));
-  }
+  };
 };
 
 struct Exp {
   template <typename T>
   T operator()(T x) {
     return fast_exp(x);
-  }
+  };
 
   complex64_t operator()(complex64_t x) {
     return std::exp(x);
   }
 };
 
-struct Expm1 {
-  template <typename T>
-  T operator()(T x) {
-    return expm1(x);
-  }
-};
-
 struct Floor {
   template <typename T>
   T operator()(T x) {
     return std::floor(x);
-  }
+  };
   int8_t operator()(int8_t x) {
     return x;
-  }
+  };
   int16_t operator()(int16_t x) {
     return x;
-  }
+  };
   int32_t operator()(int32_t x) {
     return x;
-  }
+  };
   int64_t operator()(int64_t x) {
     return x;
-  }
+  };
   uint8_t operator()(uint8_t x) {
     return x;
-  }
+  };
   uint16_t operator()(uint16_t x) {
     return x;
-  }
+  };
   uint32_t operator()(uint32_t x) {
     return x;
-  }
+  };
   uint64_t operator()(uint64_t x) {
     return x;
-  }
+  };
   bool operator()(bool x) {
     return x;
-  }
+  };
 };
 
 struct Log {
   template <typename T>
   T operator()(T x) {
     return std::log(x);
-  }
+  };
 };
 
 struct Log2 {
   template <typename T>
   T operator()(T x) {
     return std::log2(x);
-  }
+  };
 };
 
 struct Log10 {
   template <typename T>
   T operator()(T x) {
     return std::log10(x);
-  }
+  };
 };
 
 struct Log1p {
   template <typename T>
   T operator()(T x) {
     return log1p(x);
-  }
+  };
 };
 
 struct LogicalNot {
   template <typename T>
   T operator()(T x) {
     return !x;
-  }
+  };
 };
 
 struct Negative {
   template <typename T>
   T operator()(T x) {
     return -x;
-  }
+  };
 };
 
 struct Round {
@@ -379,49 +359,49 @@ struct Sin {
   template <typename T>
   T operator()(T x) {
     return std::sin(x);
-  }
+  };
 };
 
 struct Sinh {
   template <typename T>
   T operator()(T x) {
     return std::sinh(x);
-  }
+  };
 };
 
 struct Square {
   template <typename T>
   T operator()(T x) {
     return x * x;
-  }
+  };
 };
 
 struct Sqrt {
   template <typename T>
   T operator()(T x) {
     return std::sqrt(x);
-  }
+  };
 };
 
 struct Rsqrt {
   template <typename T>
   T operator()(T x) {
     return static_cast<decltype(x)>(1.0) / std::sqrt(x);
-  }
+  };
 };
 
 struct Tan {
   template <typename T>
   T operator()(T x) {
     return std::tan(x);
-  }
+  };
 };
 
 struct Tanh {
   template <typename T>
   T operator()(T x) {
     return std::tanh(x);
-  }
+  };
 };
 
 struct Add {
@@ -554,7 +534,7 @@ struct LogAddExp {
         ? maxval
         : static_cast<decltype(x)>(
               maxval + std::log1p(fast_exp(minval - maxval)));
-  }
+  };
 };
 
 struct Multiply {
@@ -602,14 +582,14 @@ struct LogicalAnd {
   template <typename T>
   T operator()(T x, T y) {
     return x && y;
-  }
+  };
 };
 
 struct LogicalOr {
   template <typename T>
   T operator()(T x, T y) {
     return x || y;
-  }
+  };
 };
 
 struct Select {
@@ -619,39 +599,4 @@ struct Select {
   }
 };
 
-struct BitwiseAnd {
-  template <typename T>
-  T operator()(T x, T y) {
-    return x & y;
-  }
-};
-
-struct BitwiseOr {
-  template <typename T>
-  T operator()(T x, T y) {
-    return x | y;
-  }
-};
-
-struct BitwiseXor {
-  template <typename T>
-  T operator()(T x, T y) {
-    return x ^ y;
-  }
-};
-
-struct LeftShift {
-  template <typename T>
-  T operator()(T x, T y) {
-    return x << y;
-  }
-};
-
-struct RightShift {
-  template <typename T>
-  T operator()(T x, T y) {
-    return x >> y;
-  }
-};
-
 } // namespace mlx::core::detail

mlx/backend/common/primitives.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <algorithm>
 #include <cassert>
@@ -8,9 +8,9 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/arange.h"
+#include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/ops.h"
-#include "mlx/backend/common/slicing.h"
 #include "mlx/backend/common/threefry.h"
 #include "mlx/backend/common/unary.h"
 #include "mlx/backend/common/utils.h"
@@ -22,7 +22,7 @@ namespace mlx::core {
 void Abs::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  if (issubdtype(in.dtype(), unsignedinteger)) {
+  if (is_unsigned(in.dtype())) {
     // No-op for unsigned types
     out.copy_shared_buffer(in);
   } else {
@@ -37,7 +37,7 @@ void Arange::eval(const std::vector<array>& inputs, array& out) {
 void ArcCos::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcCos());
   } else {
     throw std::invalid_argument(
@@ -49,7 +49,7 @@ void ArcCos::eval(const std::vector<array>& inputs, array& out) {
 void ArcCosh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcCosh());
   } else {
     throw std::invalid_argument(
@@ -61,7 +61,7 @@ void ArcCosh::eval(const std::vector<array>& inputs, array& out) {
 void ArcSin::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcSin());
   } else {
     throw std::invalid_argument(
@@ -73,7 +73,7 @@ void ArcSin::eval(const std::vector<array>& inputs, array& out) {
 void ArcSinh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcSinh());
   } else {
     throw std::invalid_argument(
@@ -85,7 +85,7 @@ void ArcSinh::eval(const std::vector<array>& inputs, array& out) {
 void ArcTan::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcTan());
   } else {
     throw std::invalid_argument(
@@ -97,7 +97,7 @@ void ArcTan::eval(const std::vector<array>& inputs, array& out) {
 void ArcTanh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::ArcTanh());
   } else {
     throw std::invalid_argument(
@@ -113,10 +113,65 @@ void AsType::eval(const std::vector<array>& inputs, array& out) {
   copy(in, out, ctype);
 }
 
+void AsStrided::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto& in = inputs[0];
+
+  if (!in.flags().row_contiguous) {
+    // Just ensuring that inputs[0] came from the ops which would ensure the
+    // input is row contiguous.
+    throw std::runtime_error(
+        "AsStrided must be used with row contiguous arrays only.");
+  }
+
+  // Compute the flags given the shape and strides
+  bool row_contiguous = true, col_contiguous = true;
+  size_t r = 1, c = 1;
+  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
+  }
+  auto flags = in.flags();
+  // TODO: Compute the contiguous flag in a better way cause now we are
+  //       unnecessarily strict.
+  flags.contiguous = row_contiguous || col_contiguous;
+  flags.row_contiguous = row_contiguous;
+  flags.col_contiguous = col_contiguous;
+
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
+
+  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
+}
+
+void Broadcast::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  if (out.size() == 0) {
+    out.set_data(nullptr);
+    return;
+  }
+  std::vector<size_t> strides(out.ndim(), 0);
+  int diff = out.ndim() - in.ndim();
+  for (int i = in.ndim() - 1; i >= 0; --i) {
+    strides[i + diff] = (in.shape()[i] == 1) ? 0 : in.strides()[i];
+  }
+  auto flags = in.flags();
+  if (out.size() > in.size()) {
+    flags.row_contiguous = flags.col_contiguous = false;
+  }
+  out.copy_shared_buffer(in, strides, flags, in.data_size());
+}
+
 void Ceil::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  if (issubdtype(in.dtype(), inexact)) {
+  if (not is_integral(in.dtype())) {
     unary_fp(in, out, detail::Ceil());
   } else {
     // No-op integer types
@@ -148,21 +203,15 @@ void Concatenate::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void Conjugate::eval(const std::vector<array>& inputs, array& out) {
+void Copy::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  const auto& in = inputs[0];
-  if (out.dtype() == complex64) {
-    unary_fp(in, out, detail::Conjugate());
-  } else {
-    throw std::invalid_argument(
-        "[conjugate] conjugate must be called on complex input.");
-  }
+  out.copy_shared_buffer(inputs[0]);
 }
 
 void Cos::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Cos());
   } else {
     throw std::invalid_argument(
@@ -174,7 +223,7 @@ void Cos::eval(const std::vector<array>& inputs, array& out) {
 void Cosh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Cosh());
   } else {
     throw std::invalid_argument(
@@ -183,6 +232,25 @@ void Cosh::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void CustomVJP::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
+       i++, j++) {
+    outputs[i].copy_shared_buffer(inputs[j]);
+  }
+}
+
+void Depends::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0; i < outputs.size(); i++) {
+    outputs[i].copy_shared_buffer(inputs[i]);
+  }
+}
+
 void Erf::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
@@ -226,7 +294,7 @@ void ErfInv::eval(const std::vector<array>& inputs, array& out) {
 void Exp::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Exp());
   } else {
     throw std::invalid_argument(
@@ -235,22 +303,10 @@ void Exp::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void Expm1::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
-    unary_fp(in, out, detail::Expm1());
-  } else {
-    throw std::invalid_argument(
-        "[expm1] Cannot exponentiate elements in array"
-        " with non floating point type.");
-  }
-}
-
 void Floor::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  if (issubdtype(in.dtype(), inexact)) {
+  if (not is_integral(in.dtype())) {
     unary_fp(in, out, detail::Floor());
   } else {
     // No-op integer types
@@ -276,7 +332,7 @@ void Full::eval(const std::vector<array>& inputs, array& out) {
 void Log::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     switch (base_) {
       case Base::e:
         unary_fp(in, out, detail::Log());
@@ -298,7 +354,7 @@ void Log::eval(const std::vector<array>& inputs, array& out) {
 void Log1p::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Log1p());
   } else {
     throw std::invalid_argument(
@@ -313,6 +369,20 @@ void LogicalNot::eval(const std::vector<array>& inputs, array& out) {
   unary(in, out, detail::LogicalNot());
 }
 
+void LogicalAnd::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2); // LogicalAnd requires two input arrays
+  auto& in1 = inputs[0];
+  auto& in2 = inputs[1];
+  binary(in1, in2, out, detail::LogicalAnd());
+}
+
+void LogicalOr::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2); // LogicalOr requires two input arrays
+  auto& in1 = inputs[0];
+  auto& in2 = inputs[1];
+  binary(in1, in2, out, detail::LogicalOr());
+}
+
 void Negative::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
@@ -401,30 +471,24 @@ void RandomBits::eval(const std::vector<array>& inputs, array& out) {
 void Reshape::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-
-  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
-
-  if (copy_necessary) {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
-    auto out_strides = make_contiguous_strides<size_t>(in.shape());
-    copy_inplace<size_t>(
-        in,
-        out,
-        in.shape(),
-        in.strides(),
-        out_strides,
-        0,
-        0,
-        CopyType::General);
+  if (in.flags().row_contiguous) {
+    // For row contiguous reshapes:
+    // - Shallow copy the buffer
+    // - If reshaping into a vector (all singleton dimensions except one) it
+    //    becomes col contiguous again.
+    auto flags = in.flags();
+    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
+    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
+    out.copy_shared_buffer(in, out.strides(), flags, in.data_size());
   } else {
-    shared_buffer_reshape(in, out_strides, out);
+    copy(in, out, in.data_size() == 1 ? CopyType::Scalar : CopyType::General);
   }
 }
 
 void Round::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  if (issubdtype(in.dtype(), inexact)) {
+  if (not is_integral(in.dtype())) {
     unary_fp(in, out, detail::Round());
   } else {
     // No-op integer types
@@ -435,7 +499,7 @@ void Round::eval(const std::vector<array>& inputs, array& out) {
 void Sigmoid::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Sigmoid());
   } else {
     throw std::invalid_argument(
@@ -457,7 +521,7 @@ void Sign::eval(const std::vector<array>& inputs, array& out) {
 void Sin::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Sin());
   } else {
     throw std::invalid_argument(
@@ -469,7 +533,7 @@ void Sin::eval(const std::vector<array>& inputs, array& out) {
 void Sinh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Sinh());
   } else {
     throw std::invalid_argument(
@@ -484,67 +548,96 @@ void Slice::eval(const std::vector<array>& inputs, array& out) {
     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, start_indices_, strides_);
-
-  // 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);
+  auto strides = in.strides();
+  auto flags = in.flags();
+  size_t data_offset = 0;
+  for (int i = 0; i < in.ndim(); ++i) {
+    data_offset += start_indices_[i] * in.strides()[i];
+    strides[i] *= strides_[i];
   }
+
+  // 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);
+    }
+  }
+
+  if (data_size == 1) {
+    // Broadcasted scalar array is contiguous.
+    flags.contiguous = true;
+  } else if (data_size == in.data_size()) {
+    // Means we sliced a broadcasted dimension so leave the "no holes" flag
+    // alone.
+  } else {
+    // We sliced something. So either we are row or col contiguous or we
+    // punched a hole.
+    flags.contiguous &= flags.row_contiguous || flags.col_contiguous;
+  }
+
+  out.copy_shared_buffer(in, strides, flags, data_size, data_offset);
 }
 
-void SliceUpdate::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  if (out.size() == 0) {
-    out.set_data(nullptr);
-    return;
-  }
+void Split::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
 
   auto& in = inputs[0];
-  auto& upd = inputs[1];
 
-  if (upd.size() == 0) {
-    out.copy_shared_buffer(in);
-    return;
+  auto compute_new_flags = [](const auto& shape,
+                              const auto& strides,
+                              size_t in_data_size,
+                              auto flags) {
+    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 = shape.size() - 1; ri >= 0; i++, ri--) {
+      flags.col_contiguous &= strides[i] == f_stride || shape[i] == 1;
+      flags.row_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
+      f_stride *= shape[i];
+      b_stride *= shape[ri];
+      if (strides[i] > 0) {
+        data_size *= shape[i];
+      }
+    }
+
+    if (data_size == 1) {
+      // Broadcasted scalar array is contiguous.
+      flags.contiguous = true;
+    } else if (data_size == in_data_size) {
+      // Means we sliced a broadcasted dimension so leave the "no holes" flag
+      // alone.
+    } else {
+      // We sliced something. So either we are row or col contiguous or we
+      // punched a hole.
+      flags.contiguous &= flags.row_contiguous || flags.col_contiguous;
+    }
+
+    return std::pair<decltype(flags), size_t>{flags, data_size};
+  };
+
+  std::vector<int> indices(1, 0);
+  indices.insert(indices.end(), indices_.begin(), indices_.end());
+  for (int i = 0; i < indices.size(); i++) {
+    size_t offset = indices[i] * in.strides()[axis_];
+    auto [new_flags, data_size] = compute_new_flags(
+        outputs[i].shape(), in.strides(), in.data_size(), in.flags());
+    outputs[i].copy_shared_buffer(
+        in, in.strides(), new_flags, data_size, offset);
   }
-
-  // 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 Square::eval(const std::vector<array>& inputs, array& out) {
@@ -563,10 +656,15 @@ void Sqrt::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+void StopGradient::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  out.copy_shared_buffer(inputs[0]);
+}
+
 void Tan::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Tan());
   } else {
     throw std::invalid_argument(
@@ -578,7 +676,7 @@ void Tan::eval(const std::vector<array>& inputs, array& out) {
 void Tanh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-  if (issubdtype(out.dtype(), inexact)) {
+  if (is_floating_point(out.dtype())) {
     unary_fp(in, out, detail::Tanh());
   } else {
     throw std::invalid_argument(
@@ -587,36 +685,38 @@ void Tanh::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-void View::eval_cpu(const std::vector<array>& inputs, array& out) {
+void Transpose::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
+  std::vector<size_t> out_strides(out.ndim());
   auto& in = inputs[0];
-  auto ibytes = size_of(in.dtype());
-  auto obytes = size_of(out.dtype());
-  // Conditions for buffer copying (disjunction):
-  // - type size is the same
-  // - type size is smaller and the last axis is contiguous
-  // - the entire array is row contiguous
-  if (ibytes == obytes || obytes < ibytes && in.strides().back() == 1 ||
-      in.flags().row_contiguous) {
-    auto strides = in.strides();
-    for (int i = 0; i < strides.size() - 1; ++i) {
-      strides[i] *= ibytes;
-      strides[i] /= obytes;
-    }
-    out.copy_shared_buffer(
-        in, strides, in.flags(), in.data_size() * obytes / ibytes);
-  } else {
-    auto tmp = array(in.shape(), in.dtype(), nullptr, {});
-    tmp.set_data(allocator::malloc_or_wait(tmp.nbytes()));
-    copy_inplace(in, tmp, CopyType::General);
-
-    auto flags = out.flags();
-    flags.contiguous = true;
-    flags.row_contiguous = true;
-    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
-    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
-    out.move_shared_buffer(tmp, out.strides(), flags, out.size());
+  for (int ax = 0; ax < axes_.size(); ++ax) {
+    out_strides[ax] = in.strides()[axes_[ax]];
   }
+
+  // Conditions for {row/col}_contiguous
+  // - array must be contiguous (no gaps)
+  // - underlying buffer size should have the same size as the array
+  // - cumulative product of shapes is equal to the strides (we can ignore axes
+  //   with size == 1)
+  //   - in the forward direction (column contiguous)
+  //   - in the reverse direction (row contiguous)
+  // - vectors are both row and col contiguous (hence if both row/col are
+  //   true, they stay true)
+  auto flags = in.flags();
+  if (flags.contiguous && in.data_size() == in.size()) {
+    size_t f_stride = 1;
+    size_t b_stride = 1;
+    flags.col_contiguous = true;
+    flags.row_contiguous = true;
+    for (int i = 0, ri = out.ndim() - 1; i < out.ndim(); ++i, --ri) {
+      flags.col_contiguous &= (out_strides[i] == f_stride || out.shape(i) == 1);
+      f_stride *= out.shape(i);
+      flags.row_contiguous &=
+          (out_strides[ri] == b_stride || out.shape(ri) == 1);
+      b_stride *= out.shape(ri);
+    }
+  }
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
 }
 
 } // namespace mlx::core

mlx/backend/common/quantized.cpp

@@ -192,7 +192,7 @@ void _qmm_dispatch_typed(
 }
 
 void _qmm_dispatch(
-    array& out,
+    array out,
     const array& x,
     const array& w,
     const array& scales,
@@ -253,81 +253,6 @@ void _qmm_dispatch(
   }
 }
 
-void _bs_qmm_dispatch(
-    array& out,
-    const array& x,
-    const array& w,
-    const array& scales,
-    const array& biases,
-    const array& lhs_indices,
-    const array& rhs_indices,
-    int bits,
-    int group_size,
-    bool transposed_w) {
-  int K = x.shape(-1);
-  int M = x.shape(-2);
-  int N = out.shape(-1);
-
-  int w_els = w.shape(-1) * w.shape(-2);
-  int g_els = scales.shape(-1) * scales.shape(-2);
-
-  const uint32_t* lhs_indices_data = lhs_indices.data<uint32_t>();
-  const uint32_t* rhs_indices_data = rhs_indices.data<uint32_t>();
-
-  for (int i = 0; i < lhs_indices.size(); i++) {
-    int x_idx = lhs_indices_data[elem_to_loc(i, lhs_indices)];
-    int w_idx = rhs_indices_data[elem_to_loc(i, rhs_indices)];
-
-    switch (x.dtype()) {
-      case float32:
-        _qmm_dispatch_typed<float>(
-            out.data<float>() + i * M * N,
-            x.data<float>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<float>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<float>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case float16:
-        _qmm_dispatch_typed<float16_t>(
-            out.data<float16_t>() + i * M * N,
-            x.data<float16_t>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<float16_t>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<float16_t>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      case bfloat16:
-        _qmm_dispatch_typed<bfloat16_t>(
-            out.data<bfloat16_t>() + i * M * N,
-            x.data<bfloat16_t>() + elem_to_loc(x_idx * M * K, x),
-            w.data<uint32_t>() + elem_to_loc(w_idx * w_els, w),
-            scales.data<bfloat16_t>() + elem_to_loc(w_idx * g_els, scales),
-            biases.data<bfloat16_t>() + elem_to_loc(w_idx * g_els, biases),
-            M,
-            N,
-            K,
-            bits,
-            group_size,
-            transposed_w);
-        break;
-      default:
-        throw std::invalid_argument(
-            "[quantized_matmul] only floating types are supported");
-    }
-  }
-}
-
 } // namespace
 
 void QuantizedMatmul::eval(const std::vector<array>& inputs, array& out) {
@@ -357,45 +282,4 @@ void QuantizedMatmul::eval(const std::vector<array>& inputs, array& out) {
   _qmm_dispatch(out, x, w, scales, biases, group_size_, bits_, transpose_);
 }
 
-void GatherQMM::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 6);
-
-  auto& x_pre = inputs[0];
-  auto& w_pre = inputs[1];
-  auto& scales_pre = inputs[2];
-  auto& biases_pre = inputs[3];
-  auto& lhs_indices = inputs[4];
-  auto& rhs_indices = inputs[5];
-
-  auto ensure_row_contiguous_last_dims = [](const array& arr) {
-    auto stride_0 = arr.strides()[arr.ndim() - 2];
-    auto stride_1 = arr.strides()[arr.ndim() - 1];
-    if (stride_0 == arr.shape(-1) && stride_1 == 1) {
-      return arr;
-    } else {
-      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
-      copy(arr, arr_copy, CopyType::General);
-      return arr_copy;
-    }
-  };
-
-  auto x = ensure_row_contiguous_last_dims(x_pre);
-  auto w = ensure_row_contiguous_last_dims(w_pre);
-  auto scales = ensure_row_contiguous_last_dims(scales_pre);
-  auto biases = ensure_row_contiguous_last_dims(biases_pre);
-
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
-  _bs_qmm_dispatch(
-      out,
-      x,
-      w,
-      scales,
-      biases,
-      lhs_indices,
-      rhs_indices,
-      group_size_,
-      bits_,
-      transpose_);
-}
-
 } // namespace mlx::core