Remove unnecessary copy from winograd

* mpi send use input as output

* move earlier

.circleci/config.yml

@@ -13,8 +13,62 @@ parameters:
   test_release:
     type: boolean
     default: false
+  linux_release:
+    type: boolean
+    default: false
 
 jobs:
+  build_documentation:
+    parameters:
+      upload-docs:
+        type: boolean
+        default: false
+    macos:
+      xcode: "15.2.0"
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install
+          command: |
+            brew install python@3.9
+            brew install doxygen
+            python3.9 -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install -r docs/requirements.txt
+            CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` pip install . -v
+      - when:
+          condition:
+            not: << parameters.upload-docs >>
+          steps:
+            - run:
+               name: Build documentation
+               command: |
+                 source env/bin/activate
+                 cd docs && doxygen && make html O=-W
+      - when:
+          condition: << parameters.upload-docs >>
+          steps:
+            - add_ssh_keys:
+                fingerprints:
+                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
+            - run:
+               name: Upload documentation
+               command: |
+                 source env/bin/activate
+                 git config user.email "mlx@group.apple.com"
+                 git config user.name "CircleCI Docs"
+                 git checkout gh-pages
+                 git rebase main
+                 cd docs
+                 git rm -rf build/html
+                 doxygen && make html O=-W
+                 git add -f build/html
+                 git commit -m "rebase"
+                 git push -f origin gh-pages
+
   linux_build_and_test:
     docker:
       - image: cimg/python:3.9
@@ -31,7 +85,7 @@ jobs:
           name: Install dependencies
           command: |
             pip install --upgrade cmake
-            pip install nanobind==2.1.0
+            pip install nanobind==2.4.0
             pip install numpy
             sudo apt-get update
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
@@ -77,13 +131,13 @@ jobs:
       - run:
           name: Install dependencies
           command: |
-            brew install python@3.8
+            brew install python@3.9
             brew install openmpi
-            python3.8 -m venv env
+            python3.9 -m venv env
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install nanobind==2.1.0
+            pip install nanobind==2.4.0
             pip install numpy
             pip install torch
             pip install tensorflow
@@ -105,7 +159,7 @@ jobs:
             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
+            mpirun --bind-to none -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: |
@@ -172,7 +226,7 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install nanobind==2.1.0
+            pip install nanobind==2.4.0
             pip install --upgrade setuptools
             pip install numpy
             pip install twine
@@ -208,7 +262,7 @@ jobs:
       - store_artifacts:
           path: dist/
 
-  build_linux_test_release:
+  build_linux_release:
     parameters:
       python_version:
         type: string
@@ -237,12 +291,13 @@ jobs:
             source env/bin/activate
             pip install --upgrade pip
             pip install --upgrade cmake
-            pip install nanobind==2.1.0
+            pip install nanobind==2.4.0
             pip install --upgrade setuptools
             pip install numpy
             pip install auditwheel
             pip install patchelf
             pip install build
+            pip install twine
             << parameters.extra_env >> \
               CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
               pip install . -v
@@ -253,6 +308,11 @@ jobs:
               python -m build --wheel
             auditwheel show dist/*
             auditwheel repair dist/* --plat manylinux_2_31_x86_64
+      - run:
+          name: Upload package
+          command: |
+            source env/bin/activate
+            twine upload wheelhouse/*
       - store_artifacts:
           path: wheelhouse/
 
@@ -272,6 +332,7 @@ workflows:
             parameters:
               xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
       - linux_build_and_test
+      - build_documentation 
 
   build_pypi_release:
     when:
@@ -288,9 +349,17 @@ workflows:
               ignore: /.*/
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
               xcode_version: ["15.0.0", "15.2.0"]
               build_env: ["PYPI_RELEASE=1"]
+      - build_documentation:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+          upload-docs: true
+
   prb:
     when:
       matches:
@@ -317,7 +386,7 @@ workflows:
       - build_release:
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
               xcode_version: ["15.0.0", "15.2.0"]
   weekly_build:
     when:
@@ -328,17 +397,17 @@ workflows:
       - build_release:
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
               xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
               build_env: ["DEV_RELEASE=1"]
   linux_test_release:
     when:
       and:
         - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.test_release >>
+        - << pipeline.parameters.linux_release >>
     jobs:
-      - build_linux_test_release:
+      - build_linux_release:
           matrix:
             parameters:
-              python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
               extra_env: ["PYPI_RELEASE=1"]

.gitignore

@@ -76,6 +76,9 @@ build/
 *.out
 *.app
 
+# Debug symbols
+*.pdb
+
 # VSCode 
 .vscode/
 .DS_Store

.pre-commit-config.yaml

@@ -1,16 +1,21 @@
 repos:
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.8
+    rev: v19.1.4
     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.8.0
+    rev: 24.10.0
     hooks:
     -   id: black
+    
 -   repo: https://github.com/pycqa/isort
     rev: 5.13.2
     hooks:
     -   id: isort
         args:
             - --profile=black
+- repo: https://github.com/cheshirekow/cmake-format-precommit
+  rev: v0.6.13
+  hooks:
+    - id: cmake-format

ACKNOWLEDGMENTS.md

@@ -7,7 +7,7 @@ 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. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`.
 - 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`.

CITATION.cff

@@ -0,0 +1,24 @@
+cff-version: 1.2.0
+title: mlx
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Awni
+    family-names: Hannun
+    affiliation: Apple
+  - given-names: Jagrit
+    family-names: Digani
+    affiliation: Apple
+  - given-names: Angelos
+    family-names: Katharopoulos
+    affiliation: Apple
+  - given-names: Ronan
+    family-names: Collobert
+    affiliation: Apple
+repository-code: 'https://github.com/ml-explore'
+abstract: >-
+  MLX: efficient and flexible machine learning on Apple
+  silicon
+license: MIT

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.24)
+cmake_minimum_required(VERSION 3.25)
 
 project(mlx LANGUAGES C CXX)
 
@@ -20,36 +20,40 @@ 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_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 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.17.3)
+  set(MLX_VERSION 0.21.1)
 endif()
+add_compile_definitions("MLX_VERSION=${MLX_VERSION}")
 
 # --------------------- Processor tests -------------------------
 
-message(STATUS "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}")
+message(
+  STATUS
+    "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}"
+)
 
-set(MLX_BUILD_ARM OFF)
-
-if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+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")
+      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()
+      set(MLX_BUILD_METAL OFF)
       message(WARNING "Building for x86_64 arch is not officially supported.")
     endif()
-    set(MLX_BUILD_METAL OFF)
-  elseif(${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64")
-    set(MLX_BUILD_ARM ON)
   endif()
 
 else()
+  set(MLX_BUILD_METAL OFF)
   message(WARNING "MLX is prioritised for Apple silicon systems using macOS.")
 endif()
 
@@ -61,207 +65,229 @@ cmake_policy(SET CMP0135 NEW)
 
 add_library(mlx)
 
-if (MLX_BUILD_METAL)
-  find_library(METAL_LIB Metal)
-  find_library(FOUNDATION_LIB Foundation)
-  find_library(QUARTZ_LIB QuartzCore)
+if(MLX_BUILD_METAL)
+  set(METAL_LIB "-framework Metal")
+  set(FOUNDATION_LIB "-framework Foundation")
+  set(QUARTZ_LIB "-framework QuartzCore")
 endif()
 
-if (MLX_BUILD_METAL AND NOT METAL_LIB)
+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)
+elseif(MLX_BUILD_METAL)
   message(STATUS "Building METAL sources")
 
-  if (MLX_METAL_DEBUG)
+  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)
+    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+    OUTPUT_VARIABLE MACOS_SDK_VERSION COMMAND_ERROR_IS_FATAL ANY)
 
-  FetchContent_Declare(
-    metal_cpp
-    URL ${METAL_CPP_URL}
-  )
+  if(${MACOS_SDK_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 macOS SDK version ${MACOS_SDK_VERSION}")
+
+  set(METAL_CPP_URL
+      https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18.zip)
+
+  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+  endif()
+  execute_process(
+    COMMAND
+      zsh "-c"
+      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
+  FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
 
   FetchContent_MakeAvailable(metal_cpp)
   target_include_directories(
-    mlx PUBLIC
-    $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
-    $<INSTALL_INTERFACE:include/metal_cpp>
-  )
-  target_link_libraries(
-    mlx PUBLIC
-    ${METAL_LIB}
-    ${FOUNDATION_LIB}
-    ${QUARTZ_LIB})
-
-  add_compile_definitions("MLX_METAL_VERSION=${MLX_METAL_VERSION}")
+    mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
+               $<INSTALL_INTERFACE:include/metal_cpp>)
+  target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
 endif()
 
-if (MLX_BUILD_CPU)
+if(WIN32)
+  if(MSVC)
+    # GGUF does not build with MSVC.
+    set(MLX_BUILD_GGUF OFF)
+    # There is no prebuilt OpenBLAS distribution for MSVC.
+    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
+  endif()
+  # Windows implementation of dlfcn.h APIs.
+  FetchContent_Declare(
+    dlfcn-win32
+    GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git
+    GIT_TAG v1.4.1
+    EXCLUDE_FROM_ALL)
+  block()
+  set(BUILD_SHARED_LIBS OFF)
+  FetchContent_MakeAvailable(dlfcn-win32)
+  endblock()
+  target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src")
+  target_link_libraries(mlx PRIVATE dl)
+endif()
+
+if(MLX_BUILD_CPU)
   find_library(ACCELERATE_LIBRARY Accelerate)
-  if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY)
+  if(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)
+  endif()
+
+  if(MLX_BUILD_ACCELERATE)
+    target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
+    add_compile_definitions(ACCELERATE_NEW_LAPACK)
+  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
+    # Download and build OpenBLAS from source code.
+    FetchContent_Declare(
+      openblas
+      GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git
+      GIT_TAG v0.3.28
+      EXCLUDE_FROM_ALL)
+    set(BUILD_STATIC_LIBS ON) # link statically
+    set(NOFORTRAN ON) # msvc has no fortran compiler
+    FetchContent_MakeAvailable(openblas)
+    target_link_libraries(mlx PRIVATE openblas)
+    target_include_directories(
+      mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include"
+                  "${CMAKE_BINARY_DIR}/generated" "${CMAKE_BINARY_DIR}")
+  else()
     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")
+      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)
+    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)
+    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.
+    target_link_libraries(mlx PRIVATE ${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)
+    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)
+    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})
+    target_link_libraries(mlx PRIVATE ${BLAS_LIBRARIES})
   endif()
 else()
   set(MLX_BUILD_ACCELERATE OFF)
 endif()
 
 find_package(MPI)
-if (MPI_FOUND)
+if(MPI_FOUND)
   execute_process(
     COMMAND zsh "-c" "mpirun --version"
     OUTPUT_VARIABLE MPI_VERSION
-    ERROR_QUIET
-  )
-  if (${MPI_VERSION} MATCHES ".*Open MPI.*")
+    ERROR_QUIET)
+  if(${MPI_VERSION} MATCHES ".*Open MPI.*")
     target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
-  elseif (MPI_VERSION STREQUAL "")
+  elseif(MPI_VERSION STREQUAL "")
     set(MPI_FOUND FALSE)
     message(
-      WARNING
-      "MPI found but mpirun is not available. Building without MPI."
-    )
+      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() 
+    message(WARNING "MPI which is not OpenMPI found. Building without MPI.")
+  endif()
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
 target_include_directories(
-  mlx
-  PUBLIC
-  $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
-  $<INSTALL_INTERFACE:include>
-)
+  mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
+             $<INSTALL_INTERFACE:include>)
 
-FetchContent_Declare(fmt
+FetchContent_Declare(
+  fmt
   GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-  GIT_TAG 10.2.1 
-  EXCLUDE_FROM_ALL
-)
+  GIT_TAG 10.2.1
+  EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(fmt)
-target_link_libraries(mlx PRIVATE fmt::fmt-header-only)
+target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>)
 
-if (MLX_BUILD_PYTHON_BINDINGS)
+if(MLX_BUILD_PYTHON_BINDINGS)
   message(STATUS "Building Python bindings.")
-  find_package(Python 3.8 COMPONENTS Interpreter Development.Module 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}")
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE nanobind_ROOT)
   find_package(nanobind CONFIG REQUIRED)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 
-if (MLX_BUILD_TESTS)
+if(MLX_BUILD_TESTS)
   include(CTest)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/tests)
 endif()
 
-if (MLX_BUILD_EXAMPLES)
+if(MLX_BUILD_EXAMPLES)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/examples/cpp)
 endif()
 
-if (MLX_BUILD_BENCHMARKS)
+if(MLX_BUILD_BENCHMARKS)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/benchmarks/cpp)
 endif()
 
-
-
 # ----------------------------- Installation -----------------------------
 include(GNUInstallDirs)
 
 # Install library
 install(
-    TARGETS mlx
-    EXPORT MLXTargets
-    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
-    INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
-)
-
+  TARGETS mlx
+  EXPORT MLXTargets
+  LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+  INCLUDES
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 
 # Install headers
 install(
-    DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/mlx
-    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
-    COMPONENT headers
-    FILES_MATCHING PATTERN "*.h"
-)
+  DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/mlx
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+  COMPONENT headers
+  FILES_MATCHING
+  PATTERN "*.h"
+  PATTERN "backend/metal/kernels.h" EXCLUDE)
 
 # Install metal dependencies
-if (MLX_BUILD_METAL)
+if(MLX_BUILD_METAL)
 
   # Install metal cpp
   install(
-      DIRECTORY ${metal_cpp_SOURCE_DIR}/
-      DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metal_cpp
-      COMPONENT metal_cpp_source
-  )
+    DIRECTORY ${metal_cpp_SOURCE_DIR}/
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metal_cpp
+    COMPONENT metal_cpp_source)
 
 endif()
 
@@ -273,31 +299,24 @@ set(MLX_CMAKE_INSTALL_MODULE_DIR share/cmake/MLX)
 install(
   EXPORT MLXTargets
   FILE MLXTargets.cmake
-  DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
-)
+  DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})
 
 include(CMakePackageConfigHelpers)
 
 write_basic_package_version_file(
   ${MLX_CMAKE_BUILD_VERSION_CONFIG}
   COMPATIBILITY SameMajorVersion
-  VERSION ${MLX_VERSION}
-)
+  VERSION ${MLX_VERSION})
 
 configure_package_config_file(
-  ${CMAKE_CURRENT_LIST_DIR}/mlx.pc.in
-  ${MLX_CMAKE_BUILD_CONFIG}
+  ${CMAKE_CURRENT_LIST_DIR}/mlx.pc.in ${MLX_CMAKE_BUILD_CONFIG}
   INSTALL_DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
   NO_CHECK_REQUIRED_COMPONENTS_MACRO
-  PATH_VARS CMAKE_INSTALL_LIBDIR CMAKE_INSTALL_INCLUDEDIR MLX_CMAKE_INSTALL_MODULE_DIR
-)
+  PATH_VARS CMAKE_INSTALL_LIBDIR CMAKE_INSTALL_INCLUDEDIR
+            MLX_CMAKE_INSTALL_MODULE_DIR)
 
-install(
-  FILES ${MLX_CMAKE_BUILD_CONFIG} ${MLX_CMAKE_BUILD_VERSION_CONFIG}
-  DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
-)
+install(FILES ${MLX_CMAKE_BUILD_CONFIG} ${MLX_CMAKE_BUILD_VERSION_CONFIG}
+        DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})
 
-install(
-  DIRECTORY ${CMAKE_MODULE_PATH}/
-  DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
-)
+install(DIRECTORY ${CMAKE_MODULE_PATH}/
+        DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})

README.md

@@ -6,7 +6,7 @@
 
 [![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
 
-MLX is an array framework for machine learning research on Apple silicon,
+MLX is an array framework for machine learning on Apple silicon,
 brought to you by Apple machine learning research.
 
 Some key features of MLX include:

benchmarks/cpp/autograd.cpp

@@ -5,35 +5,35 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"
 
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 void time_value_and_grad() {
-  auto x = ones({200, 1000});
-  eval(x);
-  auto fn = [](array x) {
+  auto x = mx::ones({200, 1000});
+  mx::eval(x);
+  auto fn = [](mx::array x) {
     for (int i = 0; i < 20; ++i) {
-      x = log(exp(x));
+      x = mx::log(mx::exp(x));
     }
-    return sum(x);
+    return mx::sum(x);
   };
 
-  auto grad_fn = grad(fn);
+  auto grad_fn = mx::grad(fn);
   auto independent_value_and_grad = [&]() {
     auto value = fn(x);
     auto dfdx = grad_fn(x);
-    return std::vector<array>{value, dfdx};
+    return std::vector<mx::array>{value, dfdx};
   };
   TIME(independent_value_and_grad);
 
-  auto value_and_grad_fn = value_and_grad(fn);
+  auto value_and_grad_fn = mx::value_and_grad(fn);
   auto combined_value_and_grad = [&]() {
     auto [value, dfdx] = value_and_grad_fn(x);
-    return std::vector<array>{value, dfdx};
+    return std::vector<mx::array>{value, dfdx};
   };
   TIME(combined_value_and_grad);
 }
 
 int main() {
-  std::cout << "Benchmarks for " << default_device() << std::endl;
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
   time_value_and_grad();
 }

benchmarks/cpp/compare_devices.cpp

@@ -4,21 +4,21 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"
 
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 void time_add_op() {
   std::vector<int> sizes(1, 1);
   for (int i = 0; i < 9; ++i) {
     sizes.push_back(10 * sizes.back());
   }
-  set_default_device(Device::cpu);
+  set_default_device(mx::Device::cpu);
   for (auto size : sizes) {
-    auto a = random::uniform({size});
-    auto b = random::uniform({size});
-    eval(a, b);
+    auto a = mx::random::uniform({size});
+    auto b = mx::random::uniform({size});
+    mx::eval(a, b);
     std::cout << "Size " << size << std::endl;
-    TIMEM("cpu", add, a, b, Device::cpu);
-    TIMEM("gpu", add, a, b, Device::gpu);
+    TIMEM("cpu", mx::add, a, b, mx::Device::cpu);
+    TIMEM("gpu", mx::add, a, b, mx::Device::gpu);
   }
 }
 

benchmarks/cpp/irregular_strides.cpp

@@ -6,105 +6,105 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"
 
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 void time_irregular_binary_ops_1D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   int size = 1000000;
   int step = 2;
-  auto a = random::uniform({size});
-  auto b = random::uniform({size});
-  eval(a, b);
+  auto a = mx::random::uniform({size});
+  auto b = mx::random::uniform({size});
+  mx::eval(a, b);
   a = slice(a, {0}, {size}, {step});
   b = slice(b, {0}, {size}, {step});
-  TIMEM("1D strided", add, a, b, device);
+  TIMEM("1D strided", mx::add, a, b, device);
 }
 
 void time_irregular_binary_ops_2D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   int size = 2048;
-  auto a = random::uniform({size, size});
-  auto b = random::uniform({size, size});
-  eval(a, b);
-  TIMEM("2D regular", add, a, b, device);
+  auto a = mx::random::uniform({size, size});
+  auto b = mx::random::uniform({size, size});
+  mx::eval(a, b);
+  TIMEM("2D regular", mx::add, a, b, device);
 
-  b = transpose(b);
-  eval(b);
-  TIMEM("2D transpose", add, a, b, device);
+  b = mx::transpose(b);
+  mx::eval(b);
+  TIMEM("2D mx::transpose", mx::add, a, b, device);
 
-  b = random::uniform({size});
-  eval(b);
-  TIMEM("2D broadcast dim 0", add, a, b, device);
+  b = mx::random::uniform({size});
+  mx::eval(b);
+  TIMEM("2D broadcast dim 0", mx::add, a, b, device);
 
-  b = reshape(b, {size, 1});
-  eval(b);
-  TIMEM("2D broadcast dim 1", add, a, b, device);
+  b = mx::reshape(b, {size, 1});
+  mx::eval(b);
+  TIMEM("2D broadcast dim 1", mx::add, a, b, device);
 }
 
 void time_irregular_binary_ops_3D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   int d0 = 32;
   int d1 = 512;
   int d2 = 512;
-  auto a = random::uniform({d0, d1, d2});
-  auto b = random::uniform({d0, d1, d2});
-  TIMEM("3D regular", add, a, b, device);
+  auto a = mx::random::uniform({d0, d1, d2});
+  auto b = mx::random::uniform({d0, d1, d2});
+  TIMEM("3D regular", mx::add, a, b, device);
 
-  b = transpose(b, {0, 2, 1});
-  TIMEM("3D transpose", add, a, b, device);
+  b = mx::transpose(b, {0, 2, 1});
+  TIMEM("3D mx::transpose", mx::add, a, b, device);
 
-  b = random::uniform({d1, d2});
-  TIMEM("3D broadcast dim 0", add, a, b, device);
+  b = mx::random::uniform({d1, d2});
+  TIMEM("3D broadcast dim 0", mx::add, a, b, device);
 
-  b = random::uniform({d0, 1, d2});
-  TIMEM("3D broadcast dim 1", add, a, b, device);
+  b = mx::random::uniform({d0, 1, d2});
+  TIMEM("3D broadcast dim 1", mx::add, a, b, device);
 
-  b = random::uniform({d0, d1, 1});
-  TIMEM("3D broadcast dim 2", add, a, b, device);
+  b = mx::random::uniform({d0, d1, 1});
+  TIMEM("3D broadcast dim 2", mx::add, a, b, device);
 
-  b = random::uniform({d2});
-  TIMEM("3D broadcast dims 0, 1", add, a, b, device);
+  b = mx::random::uniform({d2});
+  TIMEM("3D broadcast dims 0, 1", mx::add, a, b, device);
 
-  b = random::uniform({d1, 1});
-  TIMEM("3D broadcast dims 0, 2", add, a, b, device);
+  b = mx::random::uniform({d1, 1});
+  TIMEM("3D broadcast dims 0, 2", mx::add, a, b, device);
 
-  b = random::uniform({d0, 1, 1});
-  TIMEM("3D broadcast dims 1, 2", add, a, b, device);
+  b = mx::random::uniform({d0, 1, 1});
+  TIMEM("3D broadcast dims 1, 2", mx::add, a, b, device);
 }
 
 void time_irregular_binary_ops_4D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   std::vector<int> shape = {8, 8, 512, 512};
-  auto a = random::uniform(shape);
-  auto b = random::uniform(shape);
+  auto a = mx::random::uniform(shape);
+  auto b = mx::random::uniform(shape);
 
-  TIMEM("4D regular", add, a, b, device);
+  TIMEM("4D regular", mx::add, a, b, device);
 
-  b = transpose(b, {0, 1, 3, 2});
-  TIMEM("4D transpose", add, a, b, device);
+  b = mx::transpose(b, {0, 1, 3, 2});
+  TIMEM("4D mx::transpose", mx::add, a, b, device);
 
   std::string om = "4D broadcast dims ";
   for (int i = 0; i < shape.size(); ++i) {
     shape[i] = 1;
-    b = random::uniform(shape);
+    b = mx::random::uniform(shape);
     std::ostringstream msg;
     msg << om << i;
-    TIMEM(msg.str(), add, a, b, device);
+    TIMEM(msg.str(), mx::add, a, b, device);
 
     for (int j = i + 1; j < shape.size(); ++j) {
       shape[j] = 1;
       std::ostringstream msg;
       msg << om << i << ", " << j;
-      b = random::uniform(shape);
-      TIMEM(msg.str(), add, a, b, device);
+      b = mx::random::uniform(shape);
+      TIMEM(msg.str(), mx::add, a, b, device);
       shape[j] = a.shape(j);
 
       for (int k = j + 1; k < shape.size(); ++k) {
         shape[k] = 1;
         std::ostringstream msg;
         msg << om << i << ", " << j << ", " << k;
-        b = random::uniform(shape);
-        TIMEM(msg.str(), add, a, b, device);
+        b = mx::random::uniform(shape);
+        TIMEM(msg.str(), mx::add, a, b, device);
         shape[k] = a.shape(k);
       }
     }
@@ -113,83 +113,83 @@ void time_irregular_binary_ops_4D() {
 }
 
 void time_irregular_reshape() {
-  auto device = default_device();
+  auto device = mx::default_device();
   std::vector<int> shape;
-  auto reshape_fn = [&shape, device](const array& a) {
-    return reshape(a, shape, device);
+  auto reshape_fn = [&shape, device](const mx::array& a) {
+    return mx::reshape(a, shape, device);
   };
 
   int size = 64;
   int d = 2 * size;
 
-  auto a = random::uniform({d, d, d});
+  auto a = mx::random::uniform({d, d, d});
 
   shape = {8 * size, size, size};
   TIMEM("3D contiguous", reshape_fn, a);
 
-  a = transpose(a);
+  a = mx::transpose(a);
   shape = {8 * size, size, size};
-  TIMEM("3D transpose", reshape_fn, a);
+  TIMEM("3D mx::transpose", reshape_fn, a);
 
-  a = transpose(a, {1, 2, 0});
+  a = mx::transpose(a, {1, 2, 0});
   shape = {8 * size, size, size};
-  TIMEM("3D transpose dims 1 2", reshape_fn, a);
+  TIMEM("3D mx::transpose dims 1 2", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d, d}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d, d}), {d, d, d});
   TIMEM("3D broadcast dim 0", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d, 1, d}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d, 1, d}), {d, d, d});
   TIMEM("3D broadcast dim 1", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d, d, 1}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d, d, 1}), {d, d, d});
   TIMEM("3D broadcast dim 2", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d}), {d, d, d});
   TIMEM("3D broadcast dims 0, 1", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d, 1}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d, 1}), {d, d, d});
   TIMEM("3D broadcast dims 0, 2", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({d, 1, 1}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({d, 1, 1}), {d, d, d});
   TIMEM("3D broadcast dims 1, 2", reshape_fn, a);
 
-  a = broadcast_to(random::uniform({1, 1, 1}), {d, d, d});
+  a = mx::broadcast_to(mx::random::uniform({1, 1, 1}), {d, d, d});
   TIMEM("3D broadcast dims 1, 2, 3", reshape_fn, a);
 }
 
 void time_irregular_astype_1D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   int size = 1000000;
   int step = 2;
-  auto a = random::uniform({size});
+  auto a = mx::random::uniform({size});
   a = slice(a, {0}, {size}, {step});
-  TIMEM("1D strided", astype, a, int32, device);
+  TIMEM("1D strided", mx::astype, a, mx::int32, device);
 }
 
 void time_irregular_astype_2D() {
-  auto device = default_device();
+  auto device = mx::default_device();
   int size = 2048;
   std::vector<int> shape = {size, size};
 
-  auto a = random::uniform(shape);
-  TIMEM("2D regular", astype, a, int32, device);
+  auto a = mx::random::uniform(shape);
+  TIMEM("2D regular", mx::astype, a, mx::int32, device);
 
-  a = transpose(a);
-  TIMEM("2D transpose", astype, a, int32, device);
+  a = mx::transpose(a);
+  TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device);
 
-  a = broadcast_to(random::uniform({size}), shape);
-  TIMEM("2D broadcast dim 0", astype, a, int32, device);
+  a = mx::broadcast_to(mx::random::uniform({size}), shape);
+  TIMEM("2D broadcast dim 0", mx::astype, a, mx::int32, device);
 
-  a = broadcast_to(random::uniform({size, 1}), shape);
-  TIMEM("2D broadcast dim 1", astype, a, int32, device);
+  a = mx::broadcast_to(mx::random::uniform({size, 1}), shape);
+  TIMEM("2D broadcast dim 1", mx::astype, a, mx::int32, device);
 }
 
 int main(int argc, char** argv) {
   if (argc > 1) {
     bool use_gpu = !strcmp(argv[1], "gpu");
-    set_default_device(use_gpu ? Device::gpu : Device::cpu);
+    set_default_device(use_gpu ? mx::Device::gpu : mx::Device::cpu);
   }
-  std::cout << "Benchmarks for " << default_device() << std::endl;
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
   time_irregular_binary_ops_1D();
   time_irregular_binary_ops_2D();
   time_irregular_binary_ops_3D();

benchmarks/cpp/single_ops.cpp

@@ -3,20 +3,20 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"
 
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 void time_creation_ops() {
   int M = 2000;
   int N = 500;
   auto shape = {M, N};
-  auto full_fp32 = [&]() { return full(shape, 3.3f); };
+  auto full_fp32 = [&]() { return mx::full(shape, 3.3f); };
   TIME(full_fp32);
-  auto zeros_fp32 = [&]() { return zeros(shape, float32); };
+  auto zeros_fp32 = [&]() { return mx::zeros(shape, mx::float32); };
   TIME(zeros_fp32);
-  auto ones_fp32 = [&]() { return ones(shape, float32); };
+  auto ones_fp32 = [&]() { return mx::ones(shape, mx::float32); };
   TIME(ones_fp32);
 
-  auto arange_fp32 = [&]() { return arange(0.0, 10.0, 1e-4); };
+  auto arange_fp32 = [&]() { return mx::arange(0.0, 10.0, 1e-4); };
   TIME(arange_fp32);
 }
 
@@ -24,194 +24,196 @@ void time_type_conversions() {
   int M = 2000;
   int N = 500;
   auto shape = {M, N};
-  auto device = default_device();
+  auto device = mx::default_device();
 
-  auto a = zeros(shape, float32);
-  eval(a);
-  TIMEM("float32 to int32", astype, a, int32, device);
-  TIMEM("float32 to uint32", astype, a, uint32, device);
+  auto a = mx::zeros(shape, mx::float32);
+  mx::eval(a);
+  TIMEM("mx::float32 to mx::int32", mx::astype, a, mx::int32, device);
+  TIMEM("mx::float32 to mx::uint32", mx::astype, a, mx::uint32, device);
 
-  a = zeros(shape, int32);
-  eval(a);
-  TIMEM("int32 to float32", astype, a, float32, device);
+  a = mx::zeros(shape, mx::int32);
+  mx::eval(a);
+  TIMEM("mx::int32 to mx::float32", mx::astype, a, mx::float32, device);
 
-  a = zeros(shape, bool_);
-  eval(a);
-  TIMEM("bool to float32", astype, a, float32, device);
-  TIMEM("bool to int32", astype, a, int32, device);
-  TIMEM("bool to uint32", astype, a, uint32, device);
+  a = mx::zeros(shape, mx::bool_);
+  mx::eval(a);
+  TIMEM("bool to mx::float32", mx::astype, a, mx::float32, device);
+  TIMEM("bool to mx::int32", mx::astype, a, mx::int32, device);
+  TIMEM("bool to mx::uint32", mx::astype, a, mx::uint32, device);
 }
 
 void time_random_generation() {
   int M = 2000;
   int N = 500;
 
-  auto uniform = [&]() { return random::uniform({M, N}, float32); };
+  auto uniform = [&]() { return mx::random::uniform({M, N}, mx::float32); };
   TIME(uniform);
-  auto normal = [&]() { return random::normal({M, N}, float32); };
+  auto normal = [&]() { return mx::random::normal({M, N}, mx::float32); };
   TIME(normal);
 }
 
 void time_unary_ops() {
   int M = 2000;
   int N = 500;
-  auto device = default_device();
+  auto device = mx::default_device();
 
-  auto a = random::normal({M, N});
-  eval(a);
+  auto a = mx::random::normal({M, N});
+  mx::eval(a);
   TIME(mlx::core::abs, a, device);
-  TIME(negative, a, device);
-  TIME(sign, a, device);
-  TIME(square, a, device);
+  TIME(mx::negative, a, device);
+  TIME(mx::sign, a, device);
+  TIME(mx::square, a, device);
   TIME(mlx::core::sqrt, a, device);
-  TIME(rsqrt, a, device);
+  TIME(mx::rsqrt, a, device);
   TIME(mlx::core::exp, a, device);
 
-  a = random::uniform({M, N});
+  a = mx::random::uniform({M, N});
   TIME(mlx::core::log, a, device);
 }
 
 void time_binary_ops() {
   int M = 1000, N = 100, K = 10;
-  auto condition = random::randint(0, 2, {M, N, K});
-  auto a = random::uniform({M, N, K});
-  auto b = random::uniform({M, N, K});
-  auto device = default_device();
-  eval(a, b);
+  auto condition = mx::random::randint(0, 2, {M, N, K});
+  auto a = mx::random::uniform({M, N, K});
+  auto b = mx::random::uniform({M, N, K});
+  auto device = mx::default_device();
+  mx::eval(a, b);
 
-  TIME(add, a, b, device);
-  TIME(subtract, a, b, device);
-  TIME(multiply, a, b, device);
-  TIME(divide, a, b, device);
-  TIME(maximum, a, b, device);
-  TIME(minimum, a, b, device);
-  TIME(where, condition, a, b, device);
+  TIME(mx::add, a, b, device);
+  TIME(mx::subtract, a, b, device);
+  TIME(mx::multiply, a, b, device);
+  TIME(mx::divide, a, b, device);
+  TIME(mx::maximum, a, b, device);
+  TIME(mx::minimum, a, b, device);
+  TIME(mx::where, condition, a, b, device);
 
-  condition = array({true});
-  b = random::uniform({1});
-  eval(b);
-  TIMEM("scalar", add, a, b, device);
-  TIMEM("vector-scalar", subtract, a, b, device);
-  TIMEM("scalar-vector", subtract, b, a, device);
-  TIMEM("scalar", multiply, a, b, device);
-  TIMEM("vector-scalar", divide, a, b, device);
-  TIMEM("scalar-vector", divide, b, a, device);
-  TIMEM("scalar-vector", where, condition, a, b, device);
+  condition = mx::array({true});
+  b = mx::random::uniform({1});
+  mx::eval(b);
+  TIMEM("scalar", mx::add, a, b, device);
+  TIMEM("vector-scalar", mx::subtract, a, b, device);
+  TIMEM("scalar-vector", mx::subtract, b, a, device);
+  TIMEM("scalar", mx::multiply, a, b, device);
+  TIMEM("vector-scalar", mx::divide, a, b, device);
+  TIMEM("scalar-vector", mx::divide, b, a, device);
+  TIMEM("scalar-vector", mx::where, condition, a, b, device);
 
-  condition = broadcast_to(array({true}), {1000, 100});
-  a = broadcast_to(random::uniform({1}), {1000, 100});
-  b = broadcast_to(random::uniform({1}), {1000, 100});
-  eval(a, b);
-  TIMEM("scalar-scalar broadcast", add, a, b, device);
-  TIMEM("scalar-scalar broadcast", subtract, a, b, device);
-  TIMEM("scalar-scalar broadcast", multiply, a, b, device);
-  TIMEM("scalar-scalar broadcast", divide, a, b, device);
-  TIMEM("scalar-scalar broadcast", where, condition, a, b, device);
+  condition = mx::broadcast_to(mx::array({true}), {1000, 100});
+  a = mx::broadcast_to(mx::random::uniform({1}), {1000, 100});
+  b = mx::broadcast_to(mx::random::uniform({1}), {1000, 100});
+  mx::eval(a, b);
+  TIMEM("scalar-scalar broadcast", mx::add, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::subtract, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::multiply, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::divide, a, b, device);
+  TIMEM("scalar-scalar broadcast", mx::where, condition, a, b, device);
 }
 
 void time_strided_ops() {
   int M = 50, N = 50, O = 50, P = 50;
-  auto a = random::uniform({M, N, O, P});
-  auto b = random::uniform({M, N, O, P});
-  auto device = default_device();
-  eval(a, b);
-  TIMEM("non-strided", add, a, b, device);
-  a = transpose(a, {1, 0, 2, 3});
-  b = transpose(b, {3, 2, 0, 1});
-  eval(a, b);
-  TIMEM("strided", add, a, b, device);
+  auto a = mx::random::uniform({M, N, O, P});
+  auto b = mx::random::uniform({M, N, O, P});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIMEM("non-strided", mx::add, a, b, device);
+  a = mx::transpose(a, {1, 0, 2, 3});
+  b = mx::transpose(b, {3, 2, 0, 1});
+  mx::eval(a, b);
+  TIMEM("strided", mx::add, a, b, device);
 }
 
 void time_comparisons() {
   int M = 1000, N = 100, K = 10;
-  auto a = random::uniform({M, N, K});
-  auto b = random::uniform({M, N, K});
-  auto device = default_device();
-  eval(a, b);
-  TIME(equal, a, b, device);
-  TIME(greater, a, b, device);
-  TIME(greater_equal, a, b, device);
-  TIME(less, a, b, device);
-  TIME(less_equal, a, b, device);
+  auto a = mx::random::uniform({M, N, K});
+  auto b = mx::random::uniform({M, N, K});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIME(mx::equal, a, b, device);
+  TIME(mx::greater, a, b, device);
+  TIME(mx::greater_equal, a, b, device);
+  TIME(mx::less, a, b, device);
+  TIME(mx::less_equal, a, b, device);
 }
 
 void time_matvec() {
   int M = 2000, N = 200;
-  auto a = random::uniform({M, N});
-  auto b = random::uniform({N});
-  auto c = random::uniform({M});
-  eval(a, b, c);
-  auto matvec = [&]() { return matmul(a, b); };
+  auto a = mx::random::uniform({M, N});
+  auto b = mx::random::uniform({N});
+  auto c = mx::random::uniform({M});
+  mx::eval(a, b, c);
+  auto matvec = [&]() { return mx::matmul(a, b); };
   TIME(matvec);
 
-  auto matvec_transpose = [&]() { return matmul(transpose(a), c); };
+  auto matvec_transpose = [&]() { return mx::matmul(mx::transpose(a), c); };
   TIME(matvec_transpose);
 }
 
 void time_matmul() {
   int M = 1000, N = 1000, K = 1000;
-  auto a = random::uniform({M, K});
-  auto b = random::uniform({K, N});
-  auto device = default_device();
-  eval(a, b);
-  TIME(matmul, a, b, device);
+  auto a = mx::random::uniform({M, K});
+  auto b = mx::random::uniform({K, N});
+  auto device = mx::default_device();
+  mx::eval(a, b);
+  TIME(mx::matmul, a, b, device);
 
-  auto transpose_matmul = [&]() { return matmul(transpose(a), b); };
+  auto transpose_matmul = [&]() { return mx::matmul(mx::transpose(a), b); };
   TIME(transpose_matmul);
 }
 
 void time_reductions() {
-  auto a = random::normal({10000, 1000});
-  eval(a);
-  auto sum_all = [&a]() { return sum(a, false); };
+  auto a = mx::random::normal({10000, 1000});
+  mx::eval(a);
+  auto sum_all = [&a]() { return mx::sum(a, false); };
   TIME(sum_all);
 
-  auto sum_along_0 = [&a]() { return sum(a, 0, false); };
+  auto sum_along_0 = [&a]() { return mx::sum(a, 0, false); };
   TIME(sum_along_0);
 
-  auto sum_along_1 = [&a]() { return sum(a, 1, false); };
+  auto sum_along_1 = [&a]() { return mx::sum(a, 1, false); };
   TIME(sum_along_1);
 
-  auto prod_all = [&a]() { return prod(a, false); };
+  auto prod_all = [&a]() { return mx::prod(a, false); };
   TIME(prod_all);
 
-  auto all_true = [&a]() { return all(a, false); };
+  auto all_true = [&a]() { return mx::all(a, false); };
   TIME(all_true);
 
-  auto all_along_0 = [&a]() { return all(a, 0, false); };
+  auto all_along_0 = [&a]() { return mx::all(a, 0, false); };
   TIME(all_along_0);
 
-  auto all_along_1 = [&a]() { return all(a, 1, false); };
+  auto all_along_1 = [&a]() { return mx::all(a, 1, false); };
   TIME(all_along_1);
 
-  auto any_true = [&a]() { return any(a, false); };
+  auto any_true = [&a]() { return mx::any(a, false); };
   TIME(any_true);
 
-  auto argmin_along_0 = [&a]() { return argmin(a, 0, false); };
+  auto argmin_along_0 = [&a]() { return mx::argmin(a, 0, false); };
   TIME(argmin_along_0);
 
-  auto argmin_along_1 = [&a]() { return argmin(a, 1, false); };
+  auto argmin_along_1 = [&a]() { return mx::argmin(a, 1, false); };
   TIME(argmin_along_1);
 }
 
 void time_gather_scatter() {
-  auto a = random::normal({1000, 768});
-  eval(a);
-  auto indices = random::randint(0, 1000, {256});
-  eval(indices);
+  auto a = mx::random::normal({1000, 768});
+  mx::eval(a);
+  auto indices = mx::random::randint(0, 1000, {256});
+  mx::eval(indices);
 
-  auto embedding_lookup = [&a, &indices]() { return take(a, indices, 0); };
+  auto embedding_lookup = [&a, &indices]() { return mx::take(a, indices, 0); };
   TIME(embedding_lookup);
 
-  indices = random::randint(0, 768 * 1000, {256 * 768});
-  eval(indices);
+  indices = mx::random::randint(0, 768 * 1000, {256 * 768});
+  mx::eval(indices);
 
-  auto single_element_lookup = [&a, &indices]() { return take(a, indices); };
+  auto single_element_lookup = [&a, &indices]() {
+    return mx::take(a, indices);
+  };
   TIME(single_element_lookup);
 
-  indices = random::randint(0, 1000, {256});
-  auto updates = random::normal({256, 1, 768});
-  eval(indices, updates);
+  indices = mx::random::randint(0, 1000, {256});
+  auto updates = mx::random::normal({256, 1, 768});
+  mx::eval(indices, updates);
 
   auto embedding_update = [&a, &indices, &updates]() {
     return scatter(a, indices, updates, 0);
@@ -223,10 +225,10 @@ void time_gather_scatter() {
   };
   TIME(embedding_add);
 
-  a = reshape(a, {-1});
-  indices = random::randint(0, 768 * 1000, {768 * 256});
-  updates = random::normal({256 * 768, 1});
-  eval(a, indices, updates);
+  a = mx::reshape(a, {-1});
+  indices = mx::random::randint(0, 768 * 1000, {768 * 256});
+  updates = mx::random::normal({256 * 768, 1});
+  mx::eval(a, indices, updates);
 
   auto single_element_update = [&a, &indices, &updates]() {
     return scatter(a, indices, updates, 0);
@@ -240,21 +242,21 @@ void time_gather_scatter() {
 }
 
 void time_divmod() {
-  auto a = random::normal({1000});
-  auto b = random::normal({1000});
-  eval({a, b});
+  auto a = mx::random::normal({1000});
+  auto b = mx::random::normal({1000});
+  mx::eval({a, b});
 
-  auto divmod_fused = [&a, &b]() { return divmod(a, b); };
+  auto divmod_fused = [&a, &b]() { return mx::divmod(a, b); };
   TIME(divmod_fused);
 
   auto divmod_separate = [&a, &b]() {
-    return std::vector<array>{floor_divide(a, b), remainder(a, b)};
+    return std::vector<mx::array>{mx::floor_divide(a, b), mx::remainder(a, b)};
   };
   TIME(divmod_separate);
 }
 
 int main() {
-  std::cout << "Benchmarks for " << default_device() << std::endl;
+  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
   time_creation_ops();
   time_type_conversions();
   time_unary_ops();

benchmarks/python/comparative/bench_mlx.py

@@ -144,6 +144,13 @@ def reduction(op, axis, x):
     mx.eval(ys)
 
 
+def sum_and_add(axis, x, y):
+    z = x.sum(axis=axis, keepdims=True)
+    for i in range(50):
+        z = (z + y).sum(axis=axis, keepdims=True)
+    mx.eval(z)
+
+
 def softmax(axis, x):
     ys = []
     for i in range(100):
@@ -505,5 +512,8 @@ if __name__ == "__main__":
     elif args.benchmark == "selu":
         print(bench(selu, x))
 
+    elif args.benchmark == "sum_and_add":
+        print(bench(sum_and_add, axis, *xs))
+
     else:
         raise ValueError("Unknown benchmark")

benchmarks/python/conv2d_bench_cpu.py

@@ -0,0 +1,127 @@
+import argparse
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 1
+N_iter_bench = 10
+N_iter_func = 5
+mx.set_default_device(mx.cpu)
+
+
+def bench(f, a, b):
+    for i in range(N_warmup):
+        f(a, b)
+
+    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_2D(strides=(1, 1), padding=(0, 0), groups=1):
+    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)
+            ys.append(y)
+        mx.eval(ys)
+        return ys
+
+    return mx_conv_2D
+
+
+def make_pt_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
+    @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)
+            ys.append(y)
+        return ys
+
+    return pt_conv_2D
+
+
+def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, 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
+    )
+
+    a_mx = mx.array(a_np)
+    b_mx = mx.array(b_np)
+
+    a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("cpu")
+    b_pt = torch.from_numpy(b_np.transpose((0, 3, 1, 2))).to("cpu")
+
+    f_mx = make_mx_conv_2D(strides, padding, groups)
+    f_pt = make_pt_conv_2D(strides, padding, groups)
+
+    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_pt = torch.conv2d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 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, 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))}"
+        )
+
+    return time_mlx, time_torch
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run conv benchmarks")
+
+    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),
+    )
+
+    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:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch = bench_shape(
+                N, H, W, C, kH, kW, O, strides, padding, groups, 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}%"
+            )
+            if time_mlx >= 2.0 * time_torch:
+                print("ATTENTION ^^^^^^^")

benchmarks/python/conv2d_train_bench_cpu.py

@@ -0,0 +1,143 @@
+import time
+
+import mlx.core as mx
+import mlx.nn
+import mlx.optimizers as opt
+import torch
+
+
+def bench_mlx(steps: int = 20) -> float:
+    mx.set_default_device(mx.cpu)
+
+    class BenchNetMLX(mlx.nn.Module):
+        # simple encoder-decoder net
+
+        def __init__(self, in_channels, hidden_channels=32):
+            super().__init__()
+
+            self.net = mlx.nn.Sequential(
+                mlx.nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1),
+                mlx.nn.ReLU(),
+                mlx.nn.Conv2d(
+                    hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1
+                ),
+                mlx.nn.ReLU(),
+                mlx.nn.ConvTranspose2d(
+                    2 * hidden_channels, hidden_channels, kernel_size=3, padding=1
+                ),
+                mlx.nn.ReLU(),
+                mlx.nn.ConvTranspose2d(
+                    hidden_channels, in_channels, kernel_size=3, padding=1
+                ),
+            )
+
+        def __call__(self, input):
+            return self.net(input)
+
+    benchNet = BenchNetMLX(3)
+    mx.eval(benchNet.parameters())
+    optim = opt.Adam(learning_rate=1e-3)
+
+    inputs = mx.random.normal([10, 256, 256, 3])
+
+    params = benchNet.parameters()
+    optim.init(params)
+
+    state = [benchNet.state, optim.state]
+
+    def loss_fn(params, image):
+        benchNet.update(params)
+        pred_image = benchNet(image)
+        return (pred_image - image).abs().mean()
+
+    def step(params, image):
+        loss, grads = mx.value_and_grad(loss_fn)(params, image)
+        optim.update(benchNet, grads)
+        return loss
+
+    total_time = 0.0
+    print("MLX:")
+    for i in range(steps):
+        start_time = time.perf_counter()
+
+        step(benchNet.parameters(), inputs)
+        mx.eval(state)
+        end_time = time.perf_counter()
+
+        print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms")
+        total_time += (end_time - start_time) * 1000
+
+    return total_time
+
+
+def bench_torch(steps: int = 20) -> float:
+    device = torch.device("cpu")
+
+    class BenchNetTorch(torch.nn.Module):
+        # simple encoder-decoder net
+
+        def __init__(self, in_channels, hidden_channels=32):
+            super().__init__()
+
+            self.net = torch.nn.Sequential(
+                torch.nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1),
+                torch.nn.ReLU(),
+                torch.nn.Conv2d(
+                    hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1
+                ),
+                torch.nn.ReLU(),
+                torch.nn.ConvTranspose2d(
+                    2 * hidden_channels, hidden_channels, kernel_size=3, padding=1
+                ),
+                torch.nn.ReLU(),
+                torch.nn.ConvTranspose2d(
+                    hidden_channels, in_channels, kernel_size=3, padding=1
+                ),
+            )
+
+        def forward(self, input):
+            return self.net(input)
+
+    benchNet = BenchNetTorch(3).to(device)
+    optim = torch.optim.Adam(benchNet.parameters(), lr=1e-3)
+
+    inputs = torch.randn(10, 3, 256, 256, device=device)
+
+    def loss_fn(pred_image, image):
+        return (pred_image - image).abs().mean()
+
+    total_time = 0.0
+    print("PyTorch:")
+    for i in range(steps):
+        start_time = time.perf_counter()
+
+        optim.zero_grad()
+        pred_image = benchNet(inputs)
+        loss = loss_fn(pred_image, inputs)
+        loss.backward()
+        optim.step()
+
+        end_time = time.perf_counter()
+
+        print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms")
+        total_time += (end_time - start_time) * 1000
+
+    return total_time
+
+
+def main():
+    steps = 20
+    time_mlx = bench_mlx(steps)
+    time_torch = bench_torch(steps)
+
+    print(f"average time of MLX:     {time_mlx/steps:9.2f} ms")
+    print(f"total time of MLX:       {time_mlx:9.2f} ms")
+    print(f"average time of PyTorch: {time_torch/steps:9.2f} ms")
+    print(f"total time of PyTorch:   {time_torch:9.2f} ms")
+
+    diff = time_torch / time_mlx - 1.0
+    print(f"torch/mlx diff: {100. * diff:+5.2f}%")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/python/conv2d_transpose_bench_cpu.py

@@ -0,0 +1,129 @@
+import argparse
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 1
+N_iter_bench = 10
+N_iter_func = 5
+
+
+def bench(f, a, b):
+    for i in range(N_warmup):
+        f(a, b)
+
+    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_transpose_2D(strides=(1, 1), padding=(0, 0), groups=1):
+    def mx_conv_transpose_2D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = mx.conv_transpose2d(
+                a, b, stride=strides, padding=padding, groups=groups, stream=mx.cpu
+            )
+            ys.append(y)
+        mx.eval(ys)
+        return ys
+
+    return mx_conv_transpose_2D
+
+
+def make_pt_conv_transpose_2D(strides=(1, 1), padding=(0, 0), groups=1):
+    @torch.no_grad()
+    def pt_conv_transpose_2D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = torch.conv_transpose2d(
+                a, b, stride=strides, padding=padding, groups=groups
+            )
+            ys.append(y)
+        return ys
+
+    return pt_conv_transpose_2D
+
+
+def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, 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, (int(O / groups), kH, kW, C)).astype(
+        np_dtype
+    )
+
+    a_mx = mx.array(a_np)
+    b_mx = mx.array(b_np)
+
+    a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("cpu")
+    b_pt = torch.from_numpy(b_np.transpose((3, 0, 1, 2))).to("cpu")
+
+    f_mx = make_mx_conv_transpose_2D(strides, padding, groups)
+    f_pt = make_pt_conv_transpose_2D(strides, padding, groups)
+
+    time_torch = bench(f_pt, a_pt, b_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx)
+
+    out_mx = mx.conv_transpose2d(
+        a_mx, b_mx, stride=strides, padding=padding, groups=groups, stream=mx.cpu
+    )
+    out_pt = torch.conv_transpose2d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 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, 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))}"
+        )
+
+    return time_mlx, time_torch
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run conv benchmarks")
+
+    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, 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),
+    )
+
+    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:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch = bench_shape(
+                N, H, W, C, kH, kW, O, strides, padding, groups, 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}%"
+            )
+            if time_mlx >= 2.0 * time_torch:
+                print("ATTENTION ^^^^^^^")

benchmarks/python/conv3d_bench_cpu.py

@@ -0,0 +1,110 @@
+import argparse
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 1
+N_iter_bench = 10
+N_iter_func = 5
+mx.set_default_device(mx.cpu)
+
+
+def bench(f, a, b):
+    for i in range(N_warmup):
+        f(a, b)
+
+    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_3D(strides=(1, 1), padding=(0, 0), groups=1):
+    def mx_conv_3D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = mx.conv3d(a, b, stride=strides, padding=padding, groups=groups)
+            ys.append(y)
+        mx.eval(ys)
+        return ys
+
+    return mx_conv_3D
+
+
+def make_pt_conv_3D(strides=(1, 1), padding=(0, 0), groups=1):
+    @torch.no_grad()
+    def pt_conv_3D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = torch.conv3d(a, b, stride=strides, padding=padding, groups=groups)
+            ys.append(y)
+        return ys
+
+    return pt_conv_3D
+
+
+def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype):
+    scale = 1.0 / math.sqrt(kD * kH * kW * C)
+    a_np = np.random.uniform(0, 0.5, (N, D, H, W, C)).astype(np_dtype)
+    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, 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, 4, 1, 2, 3))).to("cpu")
+    b_pt = torch.from_numpy(b_np.transpose((0, 4, 1, 2, 3))).to("cpu")
+
+    f_mx = make_mx_conv_3D(strides, padding, groups)
+    f_pt = make_pt_conv_3D(strides, padding, groups)
+
+    time_torch = bench(f_pt, a_pt, b_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx)
+
+    out_mx = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    out_pt = torch.conv3d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 4, 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, D, H, W, C)}, {(O, kD, kH, kW, 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, 16, 16, 16, 16, 5, 5, 5, 16, (1, 1, 1), (2, 2, 2), 1),
+        (4, 16, 16, 16, 32, 5, 5, 5, 32, (1, 1, 1), (2, 2, 2), 1),
+    )
+
+    for dtype in dtypes:
+        print(
+            "(N,   D,   H,   W,   C), (  O, kD, kH, kW,   C),   dtype,    stride,      pads,  groups, diff%"
+        )
+        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch = bench_shape(
+                N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype
+            )
+            diff = time_torch / time_mlx - 1.0
+
+            print(
+                f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%"
+            )
+            if time_mlx >= 2.0 * time_torch:
+                print("ATTENTION ^^^^^^^")

benchmarks/python/conv3d_train_bench_cpu.py

@@ -0,0 +1,143 @@
+import time
+
+import mlx.core as mx
+import mlx.nn
+import mlx.optimizers as opt
+import torch
+
+
+def bench_mlx(steps: int = 20, shape=(10, 32, 32, 32, 3)) -> float:
+    mx.set_default_device(mx.cpu)
+
+    class BenchNetMLX(mlx.nn.Module):
+        # simple encoder-decoder net
+
+        def __init__(self, in_channels, hidden_channels=16):
+            super().__init__()
+
+            self.net = mlx.nn.Sequential(
+                mlx.nn.Conv3d(in_channels, hidden_channels, kernel_size=3, padding=1),
+                mlx.nn.ReLU(),
+                mlx.nn.Conv3d(
+                    hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1
+                ),
+                mlx.nn.ReLU(),
+                mlx.nn.ConvTranspose3d(
+                    2 * hidden_channels, hidden_channels, kernel_size=3, padding=1
+                ),
+                mlx.nn.ReLU(),
+                mlx.nn.ConvTranspose3d(
+                    hidden_channels, in_channels, kernel_size=3, padding=1
+                ),
+            )
+
+        def __call__(self, input):
+            return self.net(input)
+
+    benchNet = BenchNetMLX(3)
+    mx.eval(benchNet.parameters())
+    optim = opt.Adam(learning_rate=1e-3)
+
+    inputs = mx.random.normal(shape)
+
+    params = benchNet.parameters()
+    optim.init(params)
+
+    state = [benchNet.state, optim.state]
+
+    def loss_fn(params, image):
+        benchNet.update(params)
+        pred_image = benchNet(image)
+        return (pred_image - image).abs().mean()
+
+    def step(params, image):
+        loss, grads = mx.value_and_grad(loss_fn)(params, image)
+        optim.update(benchNet, grads)
+        return loss
+
+    total_time = 0.0
+    print("MLX:")
+    for i in range(steps):
+        start_time = time.perf_counter()
+
+        step(benchNet.parameters(), inputs)
+        mx.eval(state)
+        end_time = time.perf_counter()
+
+        print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms")
+        total_time += (end_time - start_time) * 1000
+
+    return total_time
+
+
+def bench_torch(steps: int = 20, shape=(10, 3, 32, 32, 32)) -> float:
+    device = torch.device("cpu")
+
+    class BenchNetTorch(torch.nn.Module):
+        # simple encoder-decoder net
+
+        def __init__(self, in_channels, hidden_channels=16):
+            super().__init__()
+
+            self.net = torch.nn.Sequential(
+                torch.nn.Conv3d(in_channels, hidden_channels, kernel_size=3, padding=1),
+                torch.nn.ReLU(),
+                torch.nn.Conv3d(
+                    hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1
+                ),
+                torch.nn.ReLU(),
+                torch.nn.ConvTranspose3d(
+                    2 * hidden_channels, hidden_channels, kernel_size=3, padding=1
+                ),
+                torch.nn.ReLU(),
+                torch.nn.ConvTranspose3d(
+                    hidden_channels, in_channels, kernel_size=3, padding=1
+                ),
+            )
+
+        def forward(self, input):
+            return self.net(input)
+
+    benchNet = BenchNetTorch(3).to(device)
+    optim = torch.optim.Adam(benchNet.parameters(), lr=1e-3)
+
+    inputs = torch.randn(*shape, device=device)
+
+    def loss_fn(pred_image, image):
+        return (pred_image - image).abs().mean()
+
+    total_time = 0.0
+    print("PyTorch:")
+    for i in range(steps):
+        start_time = time.perf_counter()
+
+        optim.zero_grad()
+        pred_image = benchNet(inputs)
+        loss = loss_fn(pred_image, inputs)
+        loss.backward()
+        optim.step()
+
+        end_time = time.perf_counter()
+
+        print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms")
+        total_time += (end_time - start_time) * 1000
+
+    return total_time
+
+
+def main():
+    steps = 10
+    time_mlx = bench_mlx(steps)
+    time_torch = bench_torch(steps)
+
+    print(f"average time of MLX:     {time_mlx/steps:9.2f} ms")
+    print(f"total time of MLX:       {time_mlx:9.2f} ms")
+    print(f"average time of PyTorch: {time_torch/steps:9.2f} ms")
+    print(f"total time of PyTorch:   {time_torch:9.2f} ms")
+
+    diff = time_torch / time_mlx - 1.0
+    print(f"torch/mlx diff: {100. * diff:+5.2f}%")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/python/conv3d_transpose_bench_cpu.py

@@ -0,0 +1,116 @@
+import argparse
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 1
+N_iter_bench = 10
+N_iter_func = 5
+mx.set_default_device(mx.cpu)
+
+
+def bench(f, a, b):
+    for i in range(N_warmup):
+        f(a, b)
+
+    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_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    def mx_conv_3D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = mx.conv_transpose3d(
+                a, b, stride=strides, padding=padding, groups=groups
+            )
+            ys.append(y)
+        mx.eval(ys)
+        return ys
+
+    return mx_conv_3D
+
+
+def make_pt_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    @torch.no_grad()
+    def pt_conv_3D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = torch.conv_transpose3d(
+                a, b, stride=strides, padding=padding, groups=groups
+            )
+            ys.append(y)
+        return ys
+
+    return pt_conv_3D
+
+
+def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype):
+    scale = 1.0 / math.sqrt(kD * kH * kW * C)
+    a_np = np.random.uniform(0, 0.5, (N, D, H, W, C)).astype(np_dtype)
+    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, 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, 4, 1, 2, 3))).to("cpu")
+    b_pt = torch.from_numpy(b_np.transpose((4, 0, 1, 2, 3))).to("cpu")
+
+    f_mx = make_mx_conv_3D(strides, padding, groups)
+    f_pt = make_pt_conv_3D(strides, padding, groups)
+
+    time_torch = bench(f_pt, a_pt, b_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx)
+
+    out_mx = mx.conv_transpose3d(
+        a_mx, b_mx, stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.conv_transpose3d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 4, 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, D, H, W, C)}, {(O, kD, kH, kW, 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, 16, 16, 16, 16, 5, 5, 5, 16, (1, 1, 1), (2, 2, 2), 1),
+        (4, 16, 16, 16, 32, 5, 5, 5, 32, (1, 1, 1), (2, 2, 2), 1),
+    )
+
+    for dtype in dtypes:
+        print(
+            "(N,   D,   H,   W,   C), (  O, kD, kH, kW,   C),   dtype,    stride,      pads,  groups, diff%"
+        )
+        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch = bench_shape(
+                N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype
+            )
+            diff = time_torch / time_mlx - 1.0
+
+            print(
+                f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {groups:7d}, {100. * diff:+5.2f}%"
+            )
+            if time_mlx >= 2.0 * time_torch:
+                print("ATTENTION ^^^^^^^")

benchmarks/python/scatter_bench.py

@@ -9,7 +9,7 @@ from time_utils import measure_runtime
 
 def benchmark_scatter_mlx(dst_shape, x_shape, idx_shapes):
     def scatter(dst, x, idx):
-        dst[*idx] = x
+        dst[tuple(idx)] = x
         mx.eval(dst)
 
     idx = []
@@ -23,8 +23,8 @@ def benchmark_scatter_mlx(dst_shape, x_shape, idx_shapes):
 
 
 def benchmark_scatter_torch(dst_shape, x_shape, idx_shapes, device):
-    def gather(dst, x, idx, device):
-        dst[*idx] = x
+    def scatter(dst, x, idx, device):
+        dst[tuple(idx)] = x
         if device == torch.device("mps"):
             torch.mps.synchronize()
 
@@ -34,7 +34,7 @@ def benchmark_scatter_torch(dst_shape, x_shape, idx_shapes, device):
     x = torch.randn(x_shape, dtype=torch.float32).to(device)
     dst = torch.randn(dst_shape, dtype=torch.float32).to(device)
 
-    runtime = measure_runtime(gather, dst=dst, x=x, idx=idx, device=device)
+    runtime = measure_runtime(scatter, dst=dst, x=x, idx=idx, device=device)
     print(f"PyTorch: {runtime:.3f}ms")
 
 
@@ -54,7 +54,7 @@ if __name__ == "__main__":
         (100_000, 64),
         (1_000_000, 64),
         (100_000,),
-        (2_000_00,),
+        (200_000,),
         (20_000_000,),
         (10000, 64),
         (100, 64),
@@ -91,6 +91,6 @@ if __name__ == "__main__":
 
     for dst_shape, x_shape, idx_shape in zip(dst_shapes, x_shapes, idx_shapes):
         print("=" * 20)
-        print(f"X {x_shape}, Indices {idx_shape}")
+        print(f"Dst: {dst_shape}, X {x_shape}, Indices {idx_shape}")
         benchmark_scatter_mlx(dst_shape, x_shape, idx_shape)
         benchmark_scatter_torch(dst_shape, x_shape, idx_shape, device=device)

benchmarks/python/sdpa_bench.py

@@ -1,62 +1,189 @@
+# Copyright © 2024 Apple Inc.
+
 import argparse
 import math
+import os
+import subprocess
+import time
 
 import mlx.core as mx
-from time_utils import time_fn
+import numpy as np
 
-MAX_SEQ = 300
-START_SEQ = 100
-SEQ_INCREMENT = 50
+device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+device_name = device_name.decode("utf-8").strip("\n")
+
+N_warmup = 5
+N_iter_bench = 40
+N_iter_func = 8
 
 
-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 bench(f, *args):
+    for i in range(N_warmup):
+        f(*args)
 
-        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)
+    s = time.perf_counter_ns()
+    for i in range(N_iter_bench):
+        f(*args)
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9
 
 
-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 mlx_sdpa_fused_inner(q, k, v, scale):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None)
 
-        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)
+def mlx_sdpa_unfused_inner(q, k, v, scale, f32softmax=False):
+    q_dtype = q.dtype
+    q = q * mx.array(scale, q_dtype)
+    n_q_heads = q.shape[-3]
+    n_kv_heads = k.shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    B = q.shape[0]
+    L = q.shape[2]
+
+    if n_repeats > 1:
+        q = mx.reshape(q, [B, n_kv_heads, n_repeats, L, -1])
+        k = mx.expand_dims(k, 2)
+        v = mx.expand_dims(v, 2)
+
+    scores = q @ mx.swapaxes(k, -1, -2)
+    if f32softmax:
+        scores = mx.softmax(scores.astype(mx.float32), axis=-1).astype(q_dtype)
+    else:
+        scores = mx.softmax(scores, axis=-1)
+
+    out = scores @ v
+    if n_repeats > 1:
+        out = mx.reshape(out, [B, n_q_heads, L, -1])
+
+    return out
+
+
+def mlx_spda_unfused(q, k, v, scale, transpose):
+    q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
+
+    for i in range(N_iter_func):
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_unfused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+
+    mx.eval(q_out)
+    return q_out
+
+
+def mlx_spda_fused(q, k, v, scale, transpose):
+    q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
+
+    for i in range(N_iter_func):
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_fused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+
+    mx.eval(q_out)
+    return q_out
+
+
+def bench_shape(B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose=True):
+    shape_q = (
+        (B, qsl, n_q_heads, head_dim) if transpose else (B, n_q_heads, qsl, head_dim)
+    )
+    shape_kv = (
+        (B, ksl, n_kv_heads, head_dim) if transpose else (B, n_kv_heads, ksl, head_dim)
+    )
+
+    q_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_q).astype(np_dtype)
+    k_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
+    v_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
+
+    scale = math.sqrt(1.0 / head_dim)
+
+    q_mx = mx.array(q_np)
+    k_mx = mx.array(k_np)
+    v_mx = mx.array(v_np)
+
+    time_mlx_unfused = bench(mlx_spda_unfused, q_mx, k_mx, v_mx, scale, transpose)
+    time_mlx_fused = bench(mlx_spda_fused, q_mx, k_mx, v_mx, scale, transpose)
+
+    if transpose:
+        q_mx = mx.transpose(q_mx, (0, 2, 1, 3))
+        k_mx = mx.transpose(k_mx, (0, 2, 1, 3))
+        v_mx = mx.transpose(v_mx, (0, 2, 1, 3))
+
+    o_mlx_fused = mlx_sdpa_fused_inner(q_mx, k_mx, v_mx, scale)
+    o_mlx_unfused = mlx_sdpa_unfused_inner(q_mx, k_mx, v_mx, scale, f32softmax=True)
+
+    atol = 1e-5 if np_dtype == np.float32 else 1e-4
+
+    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol):
+        print(
+            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
+        )
+
+    return time_mlx_fused, time_mlx_unfused
+
+
+def get_gflop_count(B, M, N, K):
+    return float(2.0 * N_iter_bench * N_iter_func * B * M * N * K) / float(1024.0**3)
 
 
 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)
+    parser = argparse.ArgumentParser(description="Run gemm benchmarks")
 
-    time_self_attention_sdpa()
-    time_self_attention_primitives()
+    dtypes = ("float16", "float32")[:1]
+    transposes = (False,)
+
+    # fmt: off
+    shapes_64 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,    32,    32,       64,   32,    32),
+          (  1,    64,    64,       64,   32,    32),
+          (  1,   128,   128,       64,   32,    32),
+          (  1,   256,   256,       64,   32,    32),
+          (  1,   512,   512,       64,   32,    32),
+          (  1,  1024,  1024,       64,   32,    32),
+          (  1,  2048,  2048,       64,   32,    32),
+          (  1,  4096,  4096,       64,   32,    32),
+    )
+
+    shapes_80 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,  1024,  1024,       80,   32,    32),
+          (  1,  2048,  2048,       80,   32,    32),
+          (  1,  4096,  4096,       80,   32,    32),
+    )
+
+    shapes_128 = (
+        # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
+          (  1,  1024,  1024,      128,   32,    32),
+          (  1,  2048,  2048,      128,   32,    32),
+          (  1,  4096,  4096,      128,   32,    32),
+    )
+    # fmt: on
+
+    shapes = shapes_64 + shapes_80 + shapes_128
+
+    print("  B,   qsl,   ksl, hdim, n_qh, n_kvh, tpose,   dtype, t_unfs, t_fuse, diff%")
+
+    for dtype in dtypes:
+        for transpose in transposes:
+            for B, qsl, ksl, head_dim, n_q_heads, n_kv_heads in shapes:
+                np_dtype = getattr(np, dtype)
+                time_mlx_fused, time_mlx_unfused = bench_shape(
+                    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose
+                )
+                diff = time_mlx_unfused / time_mlx_fused - 1.0
+                t_str = 1 if transpose else 0
+                print(
+                    f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:5d}, {dtype}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
+                )

benchmarks/python/sdpa_vector_bench.py

@@ -0,0 +1,58 @@
+import argparse
+import math
+
+import mlx.core as mx
+from time_utils import time_fn
+
+L = 16384
+H = 32
+H_k = H // 4
+D = 128
+dtype = mx.float16
+loops = 10
+
+
+def attention(q, k, v):
+    def _sdpa(q, k, v):
+        B, Hq, L, D = q.shape
+        _, Hk, S, _ = k.shape
+        q = q.reshape(B, Hk, Hq // Hk, L, D)
+        k = k[:, :, None, :, :]
+        v = v[:, :, None, :, :]
+        s = q @ k.transpose(0, 1, 2, 4, 3)
+        p = mx.softmax(s.astype(mx.float32), axis=-1).astype(s.dtype)
+        o = p @ v
+        return o.reshape(B, Hq, L, D)
+
+    for i in range(loops):
+        q = _sdpa(q, k, v)
+    return q
+
+
+def sdpa(q, k, v):
+    for i in range(loops):
+        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0)
+    return q
+
+
+def time_self_attention_primitives():
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    mx.eval(q, k, v)
+    time_fn(attention, q, k, v)
+
+
+def time_self_attention_sdpa():
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
+    mx.eval(q, k, v)
+    time_fn(sdpa, q, k, v)
+
+
+if __name__ == "__main__":
+    time_self_attention_sdpa()
+    time_self_attention_primitives()

cmake/extension.cmake

@@ -1,56 +1,41 @@
 include(CMakeParseArguments)
 
-###############################################################################
+# ##############################################################################
 # Build metal library
 #
 # Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib
 # from list ${SOURCES}, including list ${INCLUDE_DIRS}, depends on list ${DEPS}
 #
-# Args:
-#     TARGET: Custom target to be added for the metal library 
-#     TITLE: Name of the .metallib
-#     OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib
-#     SOURCES: List of source files
-#     INCLUDE_DIRS: List of include dirs
-#     DEPS: List of dependency files (like headers)
+# Args: TARGET: Custom target to be added for the metal library TITLE: Name of
+# the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List
+# of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
+# files (like headers)
 #
 macro(mlx_build_metallib)
   # Parse args
   set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
   set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
-  cmake_parse_arguments(
-      MTLLIB 
-      ""
-      "${oneValueArgs}"
-      "${multiValueArgs}" 
-      ${ARGN}
-  )
+  cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
   # Set output
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
-  # Collect compile options 
+  # Collect compile options
   set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math)
 
   # Prepare metallib build command
   add_custom_command(
     OUTPUT ${MTLLIB_BUILD_TARGET}
-    COMMAND xcrun -sdk macosx metal 
-                  "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
-                  ${MTLLIB_COMPILE_OPTIONS}
-                  ${MTLLIB_SOURCES}
-                  -o ${MTLLIB_BUILD_TARGET}
+    COMMAND
+      xcrun -sdk macosx metal
+      "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+      ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
     DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
     COMMAND_EXPAND_LISTS
     COMMENT "Building ${MTLLIB_TITLE}.metallib"
-    VERBATIM
-  )
+    VERBATIM)
 
   # Add metallib custom target
-  add_custom_target(
-    ${MTLLIB_TARGET}
-    DEPENDS
-    ${MTLLIB_BUILD_TARGET}
-  )
+  add_custom_target(${MTLLIB_TARGET} DEPENDS ${MTLLIB_BUILD_TARGET})
 
-endmacro(mlx_build_metallib)
+endmacro(mlx_build_metallib)

docs/src/conf.py

@@ -60,6 +60,7 @@ html_theme_options = {
     },
 }
 
+html_favicon = html_theme_options["logo"]["image_light"]
 
 # -- Options for HTMLHelp output ---------------------------------------------
 

docs/src/dev/custom_metal_kernels.rst

@@ -1,3 +1,5 @@
+.. _custom_metal_kernels:
+
 Custom Metal Kernels
 ====================
 
@@ -76,6 +78,10 @@ Putting this all together, the generated function signature for ``myexp`` is as
 
   template [[host_name("custom_kernel_myexp_float")]] [[kernel]] decltype(custom_kernel_myexp_float<float>) custom_kernel_myexp_float<float>;
 
+Note: ``grid`` and ``threadgroup`` are parameters to the Metal `dispatchThreads <https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads>`_ function.
+This means we will launch ``mx.prod(grid)`` threads, subdivided into ``threadgroup`` size threadgroups.
+For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension.
+
 Passing ``verbose=True`` to ``mx.fast.metal_kernel.__call__`` will print the generated code for debugging purposes.
 
 Using Shape/Strides

docs/src/dev/extensions.rst

@@ -420,8 +420,8 @@ element in the output.
             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 int64_t* x_strides [[buffer(6)]],
+            constant const int64_t* y_strides [[buffer(7)]],
             constant const int& ndim [[buffer(8)]],
             uint index [[thread_position_in_grid]]) {
         // Convert linear indices to offsets in array
@@ -438,24 +438,10 @@ each instantiation a unique host name so we can identify it.
 
 .. code-block:: C++
 
-    #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]]);
-
-    instantiate_axpby(float32, float);
-    instantiate_axpby(float16, half);
-    instantiate_axpby(bfloat16, bfloat16_t);
-    instantiate_axpby(complex64, complex64_t);
+    instantiate_kernel("axpby_general_float32", axpby_general, float)
+    instantiate_kernel("axpby_general_float16", axpby_general, float16_t)
+    instantiate_kernel("axpby_general_bfloat16", axpby_general, bfloat16_t)
+    instantiate_kernel("axpby_general_complex64", axpby_general, 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
@@ -494,7 +480,7 @@ below.
 
         // Prepare to encode kernel
         auto& compute_encoder = d.get_command_encoder(s.index);
-        compute_encoder->setComputePipelineState(kernel);
+        compute_encoder.set_compute_pipeline_state(kernel);
 
         // Kernel parameters are registered with buffer indices corresponding to
         // those in the kernel declaration at axpby.metal
@@ -509,14 +495,14 @@ below.
         compute_encoder.set_output_array(out, 2);
 
         // Encode alpha and beta
-        compute_encoder->setBytes(&alpha_, sizeof(float), 3);
-        compute_encoder->setBytes(&beta_, sizeof(float), 4);
+        compute_encoder.set_bytes(alpha_, 3);
+        compute_encoder.set_bytes(beta_, 4);
 
         // 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);
-        compute_encoder->setBytes(&ndim, sizeof(int), 8);
+        compute_encoder.set_vector_bytes(x.shape(), 5);
+        compute_encoder.set_vector_bytes(x.strides(), 6);
+        compute_encoder.set_bytes(y.strides(), 7);
+        compute_encoder.set_bytes(ndim, 8);
 
         // We launch 1 thread for each input and make sure that the number of
         // threads in any given threadgroup is not higher than the max allowed
@@ -530,7 +516,7 @@ below.
 
         // Launch the grid with the given number of threads divided among
         // the given threadgroups
-        compute_encoder.dispatchThreads(grid_dims, group_dims);
+        compute_encoder.dispatch_threads(grid_dims, group_dims);
     }
 
 We can now call the :meth:`axpby` operation on both the CPU and the GPU!

docs/src/dev/mlx_in_cpp.rst

@@ -0,0 +1,121 @@
+.. _mlx_in_cpp:
+
+Using MLX in C++
+================
+
+You can use MLX in a C++ project with CMake.
+
+.. note::
+
+  This guide is based one the following `example using MLX in C++ 
+  <https://github.com/ml-explore/mlx/tree/main/examples/cmake_project>`_
+
+First install MLX:
+
+.. code-block:: bash
+
+  pip install -U mlx
+
+You can also install the MLX Python package from source or just the C++
+library. For more information see the :ref:`documentation on installing MLX
+<build_and_install>`.
+
+Next make an example program in ``example.cpp``: 
+
+.. code-block:: C++
+
+  #include <iostream>
+
+  #include "mlx/mlx.h"
+
+  namespace mx = mlx::core;
+
+  int main() {
+    auto x = mx::array({1, 2, 3});
+    auto y = mx::array({1, 2, 3});
+    std::cout << x + y << std::endl;
+    return 0;
+  }
+
+The next step is to setup a CMake file in ``CMakeLists.txt``:
+
+.. code-block:: cmake
+
+  cmake_minimum_required(VERSION 3.27)
+
+  project(example LANGUAGES CXX)
+
+  set(CMAKE_CXX_STANDARD 17)
+  set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+
+Depending on how you installed MLX, you may need to tell CMake where to
+find it. 
+
+If you installed MLX with Python, then add the following to the CMake file:
+
+.. code-block:: cmake
+
+  find_package(
+    Python 3.9
+    COMPONENTS Interpreter Development.Module
+    REQUIRED)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE MLX_ROOT)
+
+If you installed the MLX C++ package to a system path, then CMake should be
+able to find it. If you installed it to a non-standard location or CMake can't
+find MLX then set ``MLX_ROOT`` to the location where MLX is installed:
+
+.. code-block:: cmake
+
+  set(MLX_ROOT "/path/to/mlx/")
+
+Next, instruct CMake to find MLX:
+
+.. code-block:: cmake
+
+  find_package(MLX CONFIG REQUIRED)
+
+Finally, add the ``example.cpp`` program as an executable and link MLX.
+
+.. code-block:: cmake
+
+  add_executable(example example.cpp)
+  target_link_libraries(example PRIVATE mlx)
+
+You can build the example with:
+
+.. code-block:: bash
+
+  cmake -B build -DCMAKE_BUILD_TYPE=Release
+  cmake --build build
+
+And run it with:
+
+.. code-block:: bash
+
+  ./build/example
+
+Note ``find_package(MLX CONFIG REQUIRED)`` sets the following variables:
+
+.. list-table:: Package Variables
+   :widths: 20 20 
+   :header-rows: 1
+
+   * - Variable 
+     - Description 
+   * - MLX_FOUND
+     - ``True`` if MLX is found
+   * - MLX_INCLUDE_DIRS
+     - Include directory
+   * - MLX_LIBRARIES
+     - Libraries to link against
+   * - MLX_CXX_FLAGS
+     - Additional compiler flags
+   * - MLX_BUILD_ACCELERATE
+     - ``True`` if MLX was built with Accelerate 
+   * - MLX_BUILD_METAL
+     - ``True`` if MLX was built with Metal

docs/src/index.rst

@@ -45,6 +45,7 @@ are the CPU and GPU.
    usage/numpy
    usage/distributed
    usage/using_streams
+   usage/export
 
 .. toctree::
    :caption: Examples
@@ -61,6 +62,7 @@ are the CPU and GPU.
    python/array
    python/data_types
    python/devices_and_streams
+   python/export
    python/ops
    python/random
    python/transforms
@@ -86,3 +88,4 @@ are the CPU and GPU.
    dev/extensions
    dev/metal_debugger
    dev/custom_metal_kernels
+   dev/mlx_in_cpp

docs/src/install.rst

@@ -1,3 +1,5 @@
+.. _build_and_install:
+
 Build and Install
 =================
 
@@ -14,7 +16,7 @@ silicon computer is
 To install from PyPI you must meet the following requirements:
 
 - Using an M series chip (Apple silicon)
-- Using a native Python >= 3.8
+- Using a native Python >= 3.9
 - macOS >= 13.5
 
 .. note::
@@ -53,7 +55,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``
+- `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0
 
 .. note::
@@ -209,7 +211,7 @@ 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.
+Metal kernel cache persists across reboots.
 
 Troubleshooting
 ^^^^^^^^^^^^^^^
@@ -240,7 +242,7 @@ x86 Shell
 
 .. _build shell:
 
-If the ouptut of ``uname -p``  is ``x86`` then your shell is running as x86 via
+If the output of ``uname -p``  is ``x86`` then your shell is running as x86 via
 Rosetta instead of natively.
 
 To fix this, find the application in Finder (``/Applications`` for iTerm,
@@ -264,4 +266,4 @@ Also check that cmake is using the correct architecture:
 
 If you see ``"x86_64"``, try re-installing ``cmake``. If you see ``"arm64"``
 but the build errors out with "Building for x86_64 on macOS is not supported."
-wipe your build cahce with ``rm -rf build/`` and try again.
+wipe your build cache with ``rm -rf build/`` and try again.

docs/src/python/data_types.rst

@@ -66,3 +66,4 @@ documentation for more information. Use :func:`issubdtype` to determine if one
    Dtype
    DtypeCategory
    issubdtype
+   finfo

docs/src/python/export.rst

@@ -0,0 +1,14 @@
+.. _export:
+
+Export Functions
+================
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+   export_function
+   import_function
+   exporter
+   export_to_dot

docs/src/python/fast.rst

@@ -12,5 +12,4 @@ Fast
   layer_norm
   rope
   scaled_dot_product_attention
-  affine_quantize
   metal_kernel

docs/src/python/linalg.rst

@@ -13,5 +13,8 @@ Linear Algebra
     norm
     cholesky
     cholesky_inv
+    cross
     qr
     svd
+    eigvalsh
+    eigh

docs/src/python/metal.rst

@@ -14,6 +14,7 @@ Metal
   get_cache_memory
   set_memory_limit
   set_cache_limit
+  set_wired_limit
   clear_cache
   start_capture
   stop_capture

docs/src/python/nn/functions.rst

@@ -13,6 +13,7 @@ simple functions.
    :template: nn-module-template.rst
 
    elu
+   celu
    gelu
    gelu_approx
    gelu_fast_approx

docs/src/python/nn/layers.rst

@@ -12,7 +12,9 @@ Layers
    ALiBi
    AvgPool1d
    AvgPool2d
+   AvgPool3d
    BatchNorm
+   CELU
    Conv1d
    Conv2d
    Conv3d
@@ -23,6 +25,7 @@ Layers
    Dropout2d
    Dropout3d
    Embedding
+   ELU
    GELU
    GLU
    GroupNorm
@@ -34,9 +37,12 @@ Layers
    LayerNorm
    LeakyReLU
    Linear
+   LogSigmoid
+   LogSoftmax
    LSTM
    MaxPool1d
    MaxPool2d
+   MaxPool3d
    Mish
    MultiHeadAttention
    PReLU
@@ -49,6 +55,7 @@ Layers
    RoPE
    SELU
    Sequential
+   Sigmoid
    SiLU
    SinusoidalPositionalEncoding
    Softmin

docs/src/python/ops.rst

@@ -80,6 +80,7 @@ Operations
    greater_equal
    hadamard_transform
    identity
+   imag
    inner
    isfinite
    isclose
@@ -88,6 +89,7 @@ Operations
    isneginf
    isposinf
    issubdtype
+   kron
    left_shift
    less
    less_equal
@@ -121,14 +123,17 @@ Operations
    pad
    power
    prod
+   put_along_axis
    quantize
    quantized_matmul
    radians
+   real
    reciprocal
    remainder
    repeat
    reshape
    right_shift
+   roll
    round
    rsqrt
    save
@@ -164,6 +169,7 @@ Operations
    tri
    tril
    triu
+   unflatten
    var
    view
    where

docs/src/python/random.rst

@@ -45,3 +45,4 @@ we use a splittable version of Threefry, which is a counter-based PRNG.
    truncated_normal
    uniform
    laplace
+   permutation

docs/src/usage/compile.rst

@@ -33,12 +33,12 @@ Let's start with a simple example:
   # Compile the function
   compiled_fun = mx.compile(fun)
 
-  # Prints: array(2.36788, dtype=float32) 
+  # Prints: array(2.36788, dtype=float32)
   print(compiled_fun(x, y))
 
 The output of both the regular function and the compiled function is the same
 up to numerical precision.
-   
+
 The first time you call a compiled function, MLX will build the compute
 graph, optimize it, and generate and compile code. This can be relatively
 slow. However, MLX will cache compiled functions, so calling a compiled
@@ -96,7 +96,7 @@ element-wise operations:
 
 .. code-block:: python
 
-  def gelu(x):  
+  def gelu(x):
       return x * (1 + mx.erf(x / math.sqrt(2))) / 2
 
 If you use this function with small arrays, it will be overhead bound. If you
@@ -136,13 +136,6 @@ Now make an array, and benchmark both functions:
 On an M1 Max the times are 15.5 and 3.1 milliseconds. The compiled ``gelu`` is
 five times faster.
 
-.. note::
-
-  As of the latest MLX, CPU functions are not fully compiled. Compiling CPU
-  functions can still be helpful, but won't typically result in as large a
-  speedup as compiling operations that run on the GPU.
-
-
 Debugging
 ---------
 
@@ -287,7 +280,7 @@ to the function. In some cases this can be pretty inconvenient. Hence,
   print(fun(mx.array(1.0)))
 
 
-Compiling Training Graphs 
+Compiling Training Graphs
 -------------------------
 
 This section will step through how to use :func:`compile` with a simple example
@@ -297,7 +290,7 @@ full forward, backward, and update with :func:`compile`.
 
 To start, here is the simple example without any compilation:
 
-.. code-block:: python 
+.. code-block:: python
 
   import mlx.core as mx
   import mlx.nn as nn
@@ -330,7 +323,7 @@ To start, here is the simple example without any compilation:
 To compile the update we can put it all in a function and compile it with the
 appropriate input and output captures. Here's the same example but compiled:
 
-.. code-block:: python 
+.. code-block:: python
 
   import mlx.core as mx
   import mlx.nn as nn
@@ -355,7 +348,7 @@ appropriate input and output captures. Here's the same example but compiled:
 
   # The state that will be captured as input and output
   state = [model.state, optimizer.state]
-      
+
   @partial(mx.compile, inputs=state, outputs=state)
   def step(x, y):
       loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
@@ -410,7 +403,7 @@ Compiling transformed functions works just as expected:
 
    In order to compile as much as possible, a transformation of a compiled
    function will not by default be compiled. To compile the transformed
-   function simply pass it through :func:`compile`. 
+   function simply pass it through :func:`compile`.
 
 You can also compile functions which themselves call compiled functions. A
 good practice is to compile the outer most function to give :func:`compile`
@@ -428,3 +421,77 @@ the most opportunity to optimize the computation graph:
   # Compiling the outer function is good to do as it will likely
   # be faster even though the inner functions are compiled
   fun = mx.compile(outer)
+
+
+
+.. _shapeless_compile:
+
+Shapeless Compilation
+---------------------
+
+When the shape of an input to a compiled function changes, the function is
+recompiled. You can compile a function once and run it on inputs with
+variable shapes by specifying ``shapeless=True`` to :func:`compile`. In this
+case changes to the shapes of the inputs do not cause the function to be
+recompiled.
+
+.. code-block:: python
+
+  def fun(x, y):
+      return mx.abs(x + y)
+
+  compiled_fun = mx.compile(fun, shapeless=True)
+
+  x = mx.array(1.0)
+  y = mx.array(-2.0)
+
+  # Firt call compiles the function
+  print(compiled_fun(x, y))
+
+  # Second call with different shapes
+  # does not recompile the function
+  x = mx.array([1.0, -6.0])
+  y = mx.array([-2.0, 3.0])
+  print(compiled_fun(x, y))
+
+
+Use shapeless compilations carefully. Since compilation is not triggered when
+shapes change, any graphs which are conditional on the input shapes will not
+work as expected. Shape-dependent computations are common and sometimes subtle
+to detect. For example:
+
+.. code-block:: python
+
+  def fun(x):
+      return x.reshape(x.shape[0] * x.shape[1], -1)
+
+  compiled_fun = mx.compile(fun, shapeless=True)
+
+  x = mx.random.uniform(shape=(2, 3, 4))
+
+  out = compiled_fun(x)
+
+  x = mx.random.uniform(shape=(5, 5, 3))
+
+  # Error, can't reshape (5, 5, 3) to (6, -1)
+  out = compiled_fun(x)
+
+The second call to the ``compiled_fun`` fails because of the call to
+:func:`reshape` which uses the static shape of ``x`` in the first call. We can
+fix this by using :func:`flatten` to avoid hardcoding the shape of ``x``:
+
+.. code-block:: python
+
+  def fun(x):
+      return x.flatten(0, 1)
+
+  compiled_fun = mx.compile(fun, shapeless=True)
+
+  x = mx.random.uniform(shape=(2, 3, 4))
+
+  out = compiled_fun(x)
+
+  x = mx.random.uniform(shape=(5, 5, 3))
+
+  # Ok
+  out = compiled_fun(x)

docs/src/usage/distributed.rst

@@ -141,12 +141,13 @@ everything else remaining the same.
     from mlx.utils import tree_map
 
     def all_reduce_grads(grads):
-        N = mx.distributed.init()
+        N = mx.distributed.init().size()
         if N == 1:
             return grads
         return tree_map(
-                lambda x: mx.distributed.all_sum(x) / N,
-                grads)
+            lambda x: mx.distributed.all_sum(x) / N,
+            grads
+        )
 
     def step(model, x, y):
         loss, grads = loss_grad_fn(model, x, y)

docs/src/usage/export.rst

@@ -0,0 +1,288 @@
+.. _export_usage:
+
+Exporting Functions
+===================
+
+.. currentmodule:: mlx.core
+
+MLX has an API to export and import functions to and from a file. This lets you
+run computations written in one MLX front-end (e.g. Python) in another MLX
+front-end (e.g. C++). 
+
+This guide walks through the basics of the MLX export API with some examples.
+To see the full list of functions check-out the :ref:`API documentation
+<export>`.
+
+Basics of Exporting 
+-------------------
+
+Let's start with a simple example:
+ 
+.. code-block:: python
+
+  def fun(x, y):
+    return x + y
+
+  x = mx.array(1.0)
+  y = mx.array(1.0)
+  mx.export_function("add.mlxfn", fun, x, y)
+
+To export a function, provide sample input arrays that the function
+can be called with. The data doesn't matter, but the shapes and types of the
+arrays do. In the above example we exported ``fun`` with two ``float32``
+scalar arrays. We can then import the function and run it:
+
+.. code-block:: python
+
+  add_fun = mx.import_function("add.mlxfn")
+
+  out, = add_fun(mx.array(1.0), mx.array(2.0))
+  # Prints: array(3, dtype=float32)
+  print(out)
+
+  out, = add_fun(mx.array(1.0), mx.array(3.0))
+  # Prints: array(4, dtype=float32)
+  print(out)
+
+  # Raises an exception
+  add_fun(mx.array(1), mx.array(3.0))
+
+  # Raises an exception
+  add_fun(mx.array([1.0, 2.0]), mx.array(3.0))
+
+Notice the third and fourth calls to ``add_fun`` raise exceptions because the
+shapes and types of the inputs are different than the shapes and types of the
+example inputs we exported the function with.
+
+Also notice that even though the original ``fun`` returns a single output
+array, the imported function always returns a tuple of one or more arrays.
+
+The inputs to :func:`export_function` and to an imported function can be
+specified as variable positional arguments or as a tuple of arrays:
+
+.. code-block:: python
+
+  def fun(x, y):
+    return x + y
+
+  x = mx.array(1.0)
+  y = mx.array(1.0)
+   
+  # Both arguments to fun are positional
+  mx.export_function("add.mlxfn", fun, x, y)
+
+  # Same as above
+  mx.export_function("add.mlxfn", fun, (x, y))
+
+  imported_fun = mx.import_function("add.mlxfn")
+
+  # Ok
+  out, = imported_fun(x, y)
+
+  # Also ok
+  out, = imported_fun((x, y))
+
+You can pass example inputs to functions as positional or keyword arguments. If
+you use keyword arguments to export the function, then you have to use the same
+keyword arguments when calling the imported function.
+
+.. code-block:: python
+
+  def fun(x, y):
+    return x + y
+
+  # One argument to fun is positional, the other is a kwarg
+  mx.export_function("add.mlxfn", fun, x, y=y)
+
+  imported_fun = mx.import_function("add.mlxfn")
+
+  # Ok
+  out, = imported_fun(x, y=y)
+
+  # Also ok
+  out, = imported_fun((x,), {"y": y})
+
+  # Raises since the keyword argument is missing
+  out, = imported_fun(x, y)
+
+  # Raises since the keyword argument has the wrong key
+  out, = imported_fun(x, z=y)
+
+
+Exporting Modules
+-----------------
+
+An :obj:`mlx.nn.Module` can be exported with or without the parameters included
+in the exported function. Here's an example:
+
+.. code-block:: python
+
+   model = nn.Linear(4, 4)
+   mx.eval(model.parameters())
+
+   def call(x):
+      return model(x)
+
+   mx.export_function("model.mlxfn", call, mx.zeros(4))
+
+In the above example, the :obj:`mlx.nn.Linear` module is exported. Its
+parameters are also saved to the ``model.mlxfn`` file.
+
+.. note::
+
+   For enclosed arrays inside an exported function, be extra careful to ensure
+   they are evaluated. The computation graph that gets exported will include
+   the computation that produces enclosed inputs.
+  
+   If the above example was missing ``mx.eval(model.parameters()``, the
+   exported function would include the random initialization of the
+   :obj:`mlx.nn.Module` parameters.
+
+If you only want to export the ``Module.__call__`` function without the
+parameters, pass them as inputs to the ``call`` wrapper:
+
+.. code-block:: python
+
+   model = nn.Linear(4, 4)
+   mx.eval(model.parameters())
+
+   def call(x, **params):
+     # Set the model's parameters to the input parameters
+     model.update(tree_unflatten(list(params.items())))
+     return model(x)
+ 
+   params = dict(tree_flatten(model.parameters()))
+   mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)
+
+
+Shapeless Exports
+-----------------
+
+Just like :func:`compile`, functions can also be exported for dynamically shaped
+inputs. Pass ``shapeless=True`` to :func:`export_function` or :func:`exporter`
+to export a function which can be used for inputs with variable shapes:
+
+.. code-block:: python
+
+  mx.export_function("fun.mlxfn", mx.abs, mx.array(0.0), shapeless=True)
+  imported_abs = mx.import_function("fun.mlxfn")
+
+  # Ok
+  out, = imported_abs(mx.array(-1.0))
+  
+  # Also ok 
+  out, = imported_abs(mx.array([-1.0, -2.0]))
+
+With ``shapeless=False`` (which is the default), the second call to
+``imported_abs`` would raise an exception with a shape mismatch.
+
+Shapeless exporting works the same as shapeless compilation and should be
+used carefully. See the :ref:`documentation on shapeless compilation
+<shapeless_compile>` for more information.
+
+Exporting Multiple Traces
+-------------------------
+
+In some cases, functions build different computation graphs for different
+input arguments. A simple way to manage this is to export to a new file with
+each set of inputs. This is a fine option in many cases. But it can be
+suboptimal if the exported functions have a large amount of duplicate constant
+data (for example the parameters of a :obj:`mlx.nn.Module`).
+
+The export API in MLX lets you export multiple traces of the same function to
+a single file by creating an exporting context manager with :func:`exporter`:
+
+.. code-block:: python
+
+  def fun(x, y=None):
+      constant = mx.array(3.0)
+      if y is not None:
+        x += y 
+      return x + constant
+
+  with mx.exporter("fun.mlxfn", fun) as exporter:
+      exporter(mx.array(1.0))
+      exporter(mx.array(1.0), y=mx.array(0.0))
+
+  imported_function = mx.import_function("fun.mlxfn")
+
+  # Call the function with y=None
+  out, = imported_function(mx.array(1.0))
+  print(out)
+
+  # Call the function with y specified
+  out, = imported_function(mx.array(1.0), y=mx.array(1.0))
+  print(out)
+
+In the above example the function constant data, (i.e. ``constant``), is only
+saved once. 
+
+Transformations with Imported Functions
+---------------------------------------
+
+Function transformations like :func:`grad`, :func:`vmap`, and :func:`compile` work
+on imported functions just like regular Python functions:
+
+.. code-block:: python
+
+  def fun(x):
+      return mx.sin(x)
+
+  x = mx.array(0.0)
+  mx.export_function("sine.mlxfn", fun, x)
+
+  imported_fun = mx.import_function("sine.mlxfn")
+
+  # Take the derivative of the imported function
+  dfdx = mx.grad(lambda x: imported_fun(x)[0])
+  # Prints: array(1, dtype=float32)
+  print(dfdx(x))
+
+  # Compile the imported function 
+  mx.compile(imported_fun)
+  # Prints: array(0, dtype=float32)
+  print(compiled_fun(x)[0])
+
+
+Importing Functions in C++
+--------------------------
+
+Importing and running functions in C++ is basically the same as importing and
+running them in Python. First, follow the :ref:`instructions <mlx_in_cpp>` to
+setup a simple C++ project that uses MLX as a library.
+
+Next, export a simple function from Python:
+
+.. code-block:: python
+
+  def fun(x, y):
+      return mx.exp(x + y)
+
+  x = mx.array(1.0)
+  y = mx.array(1.0)
+  mx.export_function("fun.mlxfn", fun, x, y)
+
+
+Import and run the function in C++ with only a few lines of code:
+
+.. code-block:: c++
+
+  auto fun = mx::import_function("fun.mlxfn");
+
+  auto inputs = {mx::array(1.0), mx::array(1.0)};
+  auto outputs = fun(inputs);
+
+  // Prints: array(2, dtype=float32)
+  std::cout << outputs[0] << std::endl;
+
+Imported functions can be transformed in C++ just like in Python. Use 
+``std::vector<mx::array>`` for positional arguments and ``std::map<std::string,
+mx::array>`` for keyword arguments when calling imported functions in C++.
+
+More Examples
+-------------
+
+Here are a few more complete examples exporting more complex functions from
+Python and importing and running them in C++:
+
+* `Inference and training a multi-layer perceptron <https://github.com/ml-explore/mlx/tree/main/examples/export>`_

docs/src/usage/function_transforms.rst

@@ -25,7 +25,7 @@ Here is a simple example:
 
 The output of :func:`grad` on :func:`sin` is simply another function. In this
 case it is the gradient of the sine function which is exactly the cosine
-function. To get the second derivative you can do: 
+function. To get the second derivative you can do:
 
 .. code-block:: shell
 
@@ -50,7 +50,7 @@ Automatic Differentiation
 .. _auto diff:
 
 Automatic differentiation in MLX works on functions rather than on implicit
-graphs. 
+graphs.
 
 .. note::
 
@@ -114,7 +114,7 @@ way to do that is the following:
 
    def loss_fn(params, x, y):
       w, b = params["weight"], params["bias"]
-      h = w * x + b 
+      h = w * x + b
       return mx.mean(mx.square(h - y))
 
    params = {"weight": mx.array(1.0), "bias": mx.array(0.0)}
@@ -132,7 +132,7 @@ way to do that is the following:
 
 Notice the tree structure of the parameters is preserved in the gradients.
 
-In some cases you may want to stop gradients from propagating through a 
+In some cases you may want to stop gradients from propagating through a
 part of the function. You can use the :func:`stop_gradient` for that.
 
 
@@ -161,19 +161,19 @@ A naive way to add the elements from two sets of vectors is with a loop:
   ys = mx.random.uniform(shape=(100, 4096))
 
   def naive_add(xs, ys):
-      return [xs[i] + ys[:, i] for i in range(xs.shape[1])]
+      return [xs[i] + ys[:, i] for i in range(xs.shape[0])]
 
 Instead you can use :func:`vmap` to automatically vectorize the addition:
 
 .. code-block:: python
-   
+
    # Vectorize over the second dimension of x and the
    # first dimension of y
-   vmap_add = mx.vmap(lambda x, y: x + y, in_axes=(1, 0))
+   vmap_add = mx.vmap(lambda x, y: x + y, in_axes=(0, 1))
 
 The ``in_axes`` parameter can be used to specify which dimensions of the
 corresponding input to vectorize over. Similarly, use ``out_axes`` to specify
-where the vectorized axes should be in the outputs. 
+where the vectorized axes should be in the outputs.
 
 Let's time these two different versions:
 
@@ -184,8 +184,8 @@ Let's time these two different versions:
   print(timeit.timeit(lambda: mx.eval(naive_add(xs, ys)), number=100))
   print(timeit.timeit(lambda: mx.eval(vmap_add(xs, ys)), number=100))
 
-On an M1 Max the naive version takes in total ``0.390`` seconds whereas the
-vectorized version takes only ``0.025`` seconds, more than ten times faster.
+On an M1 Max the naive version takes in total ``5.639`` seconds whereas the
+vectorized version takes only ``0.024`` seconds, more than 200 times faster.
 
 Of course, this operation is quite contrived. A better approach is to simply do
 ``xs + ys.T``, but for more complex functions :func:`vmap` can be quite handy.

docs/src/usage/indexing.rst

@@ -51,7 +51,7 @@ You can also use an :obj:`array` to index another :obj:`array`:
 .. code-block:: shell
 
   >>> arr = mx.arange(10)
-  >>> idx = mx.array([5, 7]) 
+  >>> idx = mx.array([5, 7])
   >>> arr[idx]
   array([5, 7], dtype=int32)
 
@@ -77,12 +77,12 @@ from the GPU. Performing bounds checking for array indices before launching the
 kernel would be extremely inefficient.
 
 Indexing with boolean masks is something that MLX may support in the future. In
-general, MLX has limited support for operations for which outputs
+general, MLX has limited support for operations for which output
 *shapes* are dependent on input *data*. Other examples of these types of
 operations which MLX does not yet support include :func:`numpy.nonzero` and the
 single input version of :func:`numpy.where`.
 
-In Place Updates 
+In Place Updates
 ----------------
 
 In place updates to indexed arrays are possible in MLX. For example:

docs/src/usage/lazy_evaluation.rst

@@ -13,7 +13,7 @@ compute graph is recorded. The actual computation only happens if an
 :func:`eval` is performed.
 
 MLX uses lazy evaluation because it has some nice features, some of which we
-describe below. 
+describe below.
 
 Transforming Compute Graphs
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -109,14 +109,14 @@ Here is a concrete example:
 
 An important behavior to be aware of is when the graph will be implicitly
 evaluated. Anytime you ``print`` an array, convert it to an
-:obj:`numpy.ndarray`, or otherwise access it's memory via :obj:`memoryview`,
+:obj:`numpy.ndarray`, or otherwise access its memory via :obj:`memoryview`,
 the graph will be evaluated. Saving arrays via :func:`save` (or any other MLX
 saving functions) will also evaluate the array.
 
 
 Calling :func:`array.item` on a scalar array will also evaluate it. In the
 example above, printing the loss (``print(loss)``) or adding the loss scalar to
-a list (``losses.append(loss.item())``) would cause a graph evaluation. If 
+a list (``losses.append(loss.item())``) would cause a graph evaluation. If
 these lines are before ``mx.eval(loss, model.parameters())`` then this
 will be a partial evaluation, computing only the forward pass.
 

docs/src/usage/numpy.rst

@@ -3,10 +3,10 @@
 Conversion to NumPy and Other Frameworks
 ========================================
 
-MLX array supports conversion between other frameworks with either:  
+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/>`_.  
+* The `Python Buffer Protocol <https://docs.python.org/3/c-api/buffer.html>`_.
+* `DLPack <https://dmlc.github.io/dlpack/latest/>`_.
 
 Let's convert an array to NumPy and back.
 
@@ -66,7 +66,7 @@ even though no in-place operations on MLX memory are executed.
 PyTorch
 -------
 
-.. warning:: 
+.. warning::
 
    PyTorch Support for :obj:`memoryview` is experimental and can break for
    multi-dimensional arrays. Casting to NumPy first is advised for now.

docs/src/usage/quick_start.rst

@@ -64,4 +64,4 @@ Other gradient transformations include :func:`vjp` for vector-Jacobian products
 and :func:`jvp` for Jacobian-vector products.
 
 Use :func:`value_and_grad` to efficiently compute both a function's output and
-gradient with respect to the function's input. 
+gradient with respect to the function's input.

docs/src/usage/saving_and_loading.rst

@@ -8,33 +8,33 @@ Saving and Loading Arrays
 MLX supports multiple array serialization formats.
 
 .. list-table:: Serialization Formats
-   :widths: 20 8 25 25 
+   :widths: 20 8 25 25
    :header-rows: 1
 
-   * - Format 
-     - Extension 
+   * - Format
+     - Extension
      - Function
-     - Notes 
-   * - NumPy 
-     - ``.npy`` 
+     - Notes
+   * - NumPy
+     - ``.npy``
      - :func:`save`
      - Single arrays only
-   * - NumPy archive 
-     - ``.npz`` 
+   * - NumPy archive
+     - ``.npz``
      - :func:`savez` and :func:`savez_compressed`
-     - Multiple arrays 
+     - Multiple arrays
    * - Safetensors
-     - ``.safetensors`` 
+     - ``.safetensors``
      - :func:`save_safetensors`
-     - Multiple arrays 
-   * - GGUF 
-     - ``.gguf`` 
+     - Multiple arrays
+   * - GGUF
+     - ``.gguf``
      - :func:`save_gguf`
      - Multiple arrays
 
 The :func:`load` function will load any of the supported serialization
 formats. It determines the format from the extensions. The output of
-:func:`load` depends on the format. 
+:func:`load` depends on the format.
 
 Here's an example of saving a single array to a file:
 

docs/src/usage/unified_memory.rst

@@ -20,7 +20,7 @@ Both ``a`` and ``b`` live in unified memory.
 
 In MLX, rather than moving arrays to devices, you specify the device when you
 run the operation. Any device can perform any operation on ``a`` and ``b``
-without needing to move them from one memory location to another. For example: 
+without needing to move them from one memory location to another. For example:
 
 .. code-block:: python
 

examples/cmake_project/CMakeLists.txt

@@ -0,0 +1,22 @@
+cmake_minimum_required(VERSION 3.27)
+
+project(example LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# Comment the following two commands only the MLX C++ library is installed and
+# set(MLX_ROOT "/path/to/mlx") directly if needed.
+find_package(
+  Python 3.9
+  COMPONENTS Interpreter Development.Module
+  REQUIRED)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE MLX_ROOT)
+
+find_package(MLX CONFIG REQUIRED)
+
+add_executable(example example.cpp)
+target_link_libraries(example PRIVATE mlx)

examples/cmake_project/README.md

@@ -0,0 +1,26 @@
+## Build and Run 
+
+Install MLX with Python:
+
+```bash
+pip install mlx>=0.22
+```
+
+Build the C++ example:
+
+```bash
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build
+```
+
+Run the C++ example:
+
+```
+./build/example
+```
+
+which should output:
+
+```
+array([2, 4, 6], dtype=int32)
+```

examples/cmake_project/example.cpp

@@ -0,0 +1,14 @@
+// Copyright © 2024 Apple Inc.
+
+#include <iostream>
+
+#include "mlx/mlx.h"
+
+namespace mx = mlx::core;
+
+int main() {
+  auto x = mx::array({1, 2, 3});
+  auto y = mx::array({1, 2, 3});
+  std::cout << x + y << std::endl;
+  return 0;
+}

examples/cpp/distributed.cpp

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

examples/cpp/linear_regression.cpp

@@ -10,7 +10,7 @@
 /**
  * An example of linear regression with MLX.
  */
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 int main() {
   int num_features = 100;
@@ -19,35 +19,35 @@ int main() {
   float learning_rate = 0.01;
 
   // True parameters
-  auto w_star = random::normal({num_features});
+  auto w_star = mx::random::normal({num_features});
 
   // The input examples (design matrix)
-  auto X = random::normal({num_examples, num_features});
+  auto X = mx::random::normal({num_examples, num_features});
 
   // Noisy labels
-  auto eps = 1e-2 * random::normal({num_examples});
-  auto y = matmul(X, w_star) + eps;
+  auto eps = 1e-2 * mx::random::normal({num_examples});
+  auto y = mx::matmul(X, w_star) + eps;
 
   // Initialize random parameters
-  array w = 1e-2 * random::normal({num_features});
+  mx::array w = 1e-2 * mx::random::normal({num_features});
 
-  auto loss_fn = [&](array w) {
-    auto yhat = matmul(X, w);
-    return (0.5f / num_examples) * sum(square(yhat - y));
+  auto loss_fn = [&](mx::array w) {
+    auto yhat = mx::matmul(X, w);
+    return (0.5f / num_examples) * mx::sum(mx::square(yhat - y));
   };
 
-  auto grad_fn = grad(loss_fn);
+  auto grad_fn = mx::grad(loss_fn);
 
   auto tic = timer::time();
   for (int it = 0; it < num_iters; ++it) {
-    auto grad = grad_fn(w);
-    w = w - learning_rate * grad;
-    eval(w);
+    auto grads = grad_fn(w);
+    w = w - learning_rate * grads;
+    mx::eval(w);
   }
   auto toc = timer::time();
 
   auto loss = loss_fn(w);
-  auto error_norm = std::sqrt(sum(square(w - w_star)).item<float>());
+  auto error_norm = std::sqrt(mx::sum(mx::square(w - w_star)).item<float>());
   auto throughput = num_iters / timer::seconds(toc - tic);
   std::cout << "Loss " << loss << ", |w - w*| = " << error_norm
             << ", Throughput " << throughput << " (it/s)." << std::endl;

examples/cpp/logistic_regression.cpp

@@ -10,7 +10,7 @@
 /**
  * An example of logistic regression with MLX.
  */
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 int main() {
   int num_features = 100;
@@ -19,35 +19,35 @@ int main() {
   float learning_rate = 0.1;
 
   // True parameters
-  auto w_star = random::normal({num_features});
+  auto w_star = mx::random::normal({num_features});
 
   // The input examples
-  auto X = random::normal({num_examples, num_features});
+  auto X = mx::random::normal({num_examples, num_features});
 
   // Labels
-  auto y = matmul(X, w_star) > 0;
+  auto y = mx::matmul(X, w_star) > 0;
 
   // Initialize random parameters
-  array w = 1e-2 * random::normal({num_features});
+  mx::array w = 1e-2 * mx::random::normal({num_features});
 
-  auto loss_fn = [&](array w) {
-    auto logits = matmul(X, w);
+  auto loss_fn = [&](mx::array w) {
+    auto logits = mx::matmul(X, w);
     auto scale = (1.0f / num_examples);
-    return scale * sum(logaddexp(array(0.0f), logits) - y * logits);
+    return scale * mx::sum(mx::logaddexp(mx::array(0.0f), logits) - y * logits);
   };
 
-  auto grad_fn = grad(loss_fn);
+  auto grad_fn = mx::grad(loss_fn);
 
   auto tic = timer::time();
   for (int it = 0; it < num_iters; ++it) {
-    auto grad = grad_fn(w);
-    w = w - learning_rate * grad;
-    eval(w);
+    auto grads = grad_fn(w);
+    w = w - learning_rate * grads;
+    mx::eval(w);
   }
   auto toc = timer::time();
 
   auto loss = loss_fn(w);
-  auto acc = sum((matmul(X, w) > 0) == y) / num_examples;
+  auto acc = mx::sum((mx::matmul(X, w) > 0) == y) / num_examples;
   auto throughput = num_iters / timer::seconds(toc - tic);
   std::cout << "Loss " << loss << ", Accuracy, " << acc << ", Throughput "
             << throughput << " (it/s)." << std::endl;

examples/cpp/metal_capture.cpp

@@ -5,27 +5,27 @@
 
 #include "mlx/mlx.h"
 
-using namespace mlx::core;
+namespace mx = 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");
+  mx::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 s2 = new_stream(mx::Device::gpu);
+  auto s3 = new_stream(mx::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);
+  auto a = mx::arange(1.f, 10.f, 1.f, mx::float32, s2);
+  auto b = mx::arange(1.f, 10.f, 1.f, mx::float32, s3);
+  auto x = mx::add(a, a, s2);
+  auto y = mx::add(b, b, s3);
 
   // The multiply will happen on the default stream.
-  std::cout << multiply(x, y) << std::endl;
+  std::cout << mx::multiply(x, y) << std::endl;
 
-  metal::stop_capture();
+  mx::metal::stop_capture();
 }

examples/cpp/tutorial.cpp

@@ -5,11 +5,11 @@
 
 #include "mlx/mlx.h"
 
-using namespace mlx::core;
+namespace mx = mlx::core;
 
 void array_basics() {
   // Make a scalar array:
-  array x(1.0);
+  mx::array x(1.0);
 
   // Get the value out of it:
   auto s = x.item<float>();
@@ -29,31 +29,31 @@ void array_basics() {
 
   // The datatype should be float32:
   auto dtype = x.dtype();
-  assert(dtype == float32);
+  assert(dtype == mx::float32);
 
   // Specify the dtype when constructing the array:
-  x = array(1, int32);
-  assert(x.dtype() == int32);
+  x = mx::array(1, mx::int32);
+  assert(x.dtype() == mx::int32);
   x.item<int>(); // OK
   // x.item<float>();  // Undefined!
 
   // Make a multidimensional array:
-  x = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
+  x = mx::array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
   // mlx is row-major by default so the first row of this array
   // is [1.0, 2.0] and the second row is [3.0, 4.0]
 
   // Make an array of shape {2, 2} filled with ones:
-  auto y = ones({2, 2});
+  auto y = mx::ones({2, 2});
 
   // Pointwise add x and y:
-  auto z = add(x, y);
+  auto z = mx::add(x, y);
 
   // Same thing:
   z = x + y;
 
   // mlx is lazy by default. At this point `z` only
   // has a shape and a type but no actual data:
-  assert(z.dtype() == float32);
+  assert(z.dtype() == mx::float32);
   assert(z.shape(0) == 2);
   assert(z.shape(1) == 2);
 
@@ -63,33 +63,33 @@ void array_basics() {
   // and inputs. When `eval` is called on an array (or arrays), the array and
   // all of its dependencies are recursively evaluated to produce the result.
   // Once an array is evaluated, it has data and is detached from its inputs.
-  eval(z);
+  mx::eval(z);
 
-  // Of course the array can still be an input to other operations. You can even
-  // call eval on the array again, this will just be a no-op:
-  eval(z); // no-op
+  // Of course the array can still be an input to other operations. You can
+  // even call eval on the array again, this will just be a no-op:
+  mx::eval(z); // no-op
 
   // Some functions or methods on arrays implicitly evaluate them. For example
   // accessing a value in an array or printing the array implicitly evaluate it:
-  z = ones({1});
+  z = mx::ones({1});
   z.item<float>(); // implicit evaluation
 
-  z = ones({2, 2});
+  z = mx::ones({2, 2});
   std::cout << z << std::endl; // implicit evaluation
 }
 
 void automatic_differentiation() {
-  auto fn = [](array x) { return square(x); };
+  auto fn = [](mx::array x) { return mx::square(x); };
 
   // Computing the derivative function of a function
-  auto grad_fn = grad(fn);
+  auto grad_fn = mx::grad(fn);
   // Call grad_fn on the input to get the derivative
-  auto x = array(1.5);
+  auto x = mx::array(1.5);
   auto dfdx = grad_fn(x);
   // dfdx is 2 * x
 
   // Get the second derivative by composing grad with grad
-  auto d2fdx2 = grad(grad(fn))(x);
+  auto d2fdx2 = mx::grad(mx::grad(fn))(x);
   // d2fdx2 is 2
 }
 

examples/export/CMakeLists.txt

@@ -0,0 +1,22 @@
+cmake_minimum_required(VERSION 3.27)
+
+project(import_mlx LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+find_package(
+  Python 3.9
+  COMPONENTS Interpreter Development.Module
+  REQUIRED)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE MLX_ROOT)
+find_package(MLX CONFIG REQUIRED)
+
+add_executable(eval_mlp eval_mlp.cpp)
+target_link_libraries(eval_mlp PRIVATE mlx)
+
+add_executable(train_mlp train_mlp.cpp)
+target_link_libraries(train_mlp PRIVATE mlx)

examples/export/README.md

@@ -0,0 +1,49 @@
+## Setup
+
+Install MLX:
+
+```bash
+pip install mlx>=0.22
+```
+
+Build the C++ examples:
+
+```bash
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build
+```
+
+## Run
+
+### Eval MLP
+
+Run the Python script to export the eval function:
+
+```bash
+python eval_mlp.py
+```
+
+Then run the C++ program to import and run the function:
+
+```
+./build/eval_mlp
+```
+
+The Python and C++ programs should output the same result.
+
+### Train MLP
+
+Run the Python script to export the model initialization and training
+functions:
+
+```bash
+python train_mlp.py
+```
+
+Then run the C++ program to import and run the functions:
+
+```
+./build/train_mlp
+```
+
+The Python and C++ programs should output the same results.

examples/export/eval_mlp.cpp

@@ -0,0 +1,25 @@
+// Copyright © 2024 Apple Inc.
+
+#include <mlx/mlx.h>
+#include <iostream>
+
+namespace mx = mlx::core;
+
+int main() {
+  int batch_size = 8;
+  int input_dim = 32;
+
+  // Make the input
+  mx::random::seed(42);
+  auto example_x = mx::random::uniform({batch_size, input_dim});
+
+  // Import the function
+  auto forward = mx::import_function("eval_mlp.mlxfn");
+
+  // Call the imported function
+  auto out = forward({example_x})[0];
+
+  std::cout << out << std::endl;
+
+  return 0;
+}

examples/export/eval_mlp.py

@@ -0,0 +1,52 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+
+
+class MLP(nn.Module):
+    """A simple MLP."""
+
+    def __init__(
+        self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
+    ):
+        super().__init__()
+        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
+        self.layers = [
+            nn.Linear(idim, odim)
+            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
+        ]
+
+    def __call__(self, x):
+        for l in self.layers[:-1]:
+            x = nn.relu(l(x))
+        return self.layers[-1](x)
+
+
+if __name__ == "__main__":
+
+    batch_size = 8
+    input_dim = 32
+    output_dim = 10
+
+    # Load the model
+    mx.random.seed(0)  # Seed for params
+    model = MLP(num_layers=5, input_dim=input_dim, hidden_dim=64, output_dim=output_dim)
+    mx.eval(model)
+
+    # Note, the model parameters are saved in the export function
+    def forward(x):
+        return model(x)
+
+    mx.random.seed(42)  # Seed for input
+    example_x = mx.random.uniform(shape=(batch_size, input_dim))
+
+    mx.export_function("eval_mlp.mlxfn", forward, example_x)
+
+    # Import in Python
+    imported_forward = mx.import_function("eval_mlp.mlxfn")
+    expected = forward(example_x)
+    (out,) = imported_forward(example_x)
+    assert mx.allclose(expected, out)
+    print(out)

examples/export/train_mlp.cpp

@@ -0,0 +1,35 @@
+// Copyright © 2024 Apple Inc.
+
+#include <mlx/mlx.h>
+#include <iostream>
+
+namespace mx = mlx::core;
+
+int main() {
+  int batch_size = 8;
+  int input_dim = 32;
+  int output_dim = 10;
+
+  auto state = mx::import_function("init_mlp.mlxfn")({});
+
+  // Make the input
+  mx::random::seed(42);
+  auto example_X = mx::random::normal({batch_size, input_dim});
+  auto example_y = mx::random::randint(0, output_dim, {batch_size});
+
+  // Import the function
+  auto step = mx::import_function("train_mlp.mlxfn");
+
+  // Call the imported function
+  for (int it = 0; it < 100; ++it) {
+    state.insert(state.end(), {example_X, example_y});
+    state = step(state);
+    eval(state);
+    auto loss = state.back();
+    state.pop_back();
+    if (it % 10 == 0) {
+      std::cout << "Loss " << loss.item<float>() << std::endl;
+    }
+  }
+  return 0;
+}

examples/export/train_mlp.py

@@ -0,0 +1,76 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+import mlx.utils
+
+
+class MLP(nn.Module):
+    """A simple MLP."""
+
+    def __init__(
+        self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
+    ):
+        super().__init__()
+        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
+        self.layers = [
+            nn.Linear(idim, odim)
+            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
+        ]
+
+    def __call__(self, x):
+        for l in self.layers[:-1]:
+            x = nn.relu(l(x))
+        return self.layers[-1](x)
+
+
+if __name__ == "__main__":
+
+    batch_size = 8
+    input_dim = 32
+    output_dim = 10
+
+    def init():
+        # Seed for the parameter initialization
+        mx.random.seed(0)
+        model = MLP(
+            num_layers=3, input_dim=input_dim, hidden_dim=64, output_dim=output_dim
+        )
+        optimizer = optim.SGD(learning_rate=1e-1)
+        optimizer.init(model.parameters())
+        state = [model.parameters(), optimizer.state]
+        tree_structure, state = zip(*mlx.utils.tree_flatten(state))
+        return model, optimizer, tree_structure, state
+
+    # Export the model parameter initialization
+    model, optimizer, tree_structure, state = init()
+    mx.eval(state)
+    mx.export_function("init_mlp.mlxfn", lambda: init()[-1])
+
+    def loss_fn(params, X, y):
+        model.update(params)
+        return nn.losses.cross_entropy(model(X), y, reduction="mean")
+
+    def step(*inputs):
+        *state, X, y = inputs
+        params, opt_state = mlx.utils.tree_unflatten(list(zip(tree_structure, state)))
+        optimizer.state = opt_state
+        loss, grads = mx.value_and_grad(loss_fn)(params, X, y)
+        params = optimizer.apply_gradients(grads, params)
+        _, state = zip(*mlx.utils.tree_flatten([params, optimizer.state]))
+        return *state, loss
+
+    # Make some random data
+    mx.random.seed(42)
+    example_X = mx.random.normal(shape=(batch_size, input_dim))
+    example_y = mx.random.randint(low=0, high=output_dim, shape=(batch_size,))
+    mx.export_function("train_mlp.mlxfn", step, *state, example_X, example_y)
+
+    # Export one step of SGD
+    imported_step = mx.import_function("train_mlp.mlxfn")
+
+    for it in range(100):
+        *state, loss = imported_step(*state, example_X, example_y)
+        if it % 10 == 0:
+            print(f"Loss {loss.item():.6}")

examples/extensions/CMakeLists.txt

@@ -11,11 +11,14 @@ 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)
+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}")
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE nanobind_ROOT)
 find_package(nanobind CONFIG REQUIRED)
 
 # ----------------------------- Extensions -----------------------------
@@ -24,16 +27,10 @@ find_package(nanobind CONFIG REQUIRED)
 add_library(mlx_ext)
 
 # Add sources
-target_sources(
-  mlx_ext
-  PUBLIC
-  ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp
-)
+target_sources(mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp)
 
 # Add include headers
-target_include_directories(
-  mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}
-)
+target_include_directories(mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR})
 
 # Link to mlx
 target_link_libraries(mlx_ext PUBLIC mlx)
@@ -43,27 +40,32 @@ target_link_libraries(mlx_ext PUBLIC mlx)
 # Build metallib
 if(MLX_BUILD_METAL)
   mlx_build_metallib(
-    TARGET mlx_ext_metallib
-    TITLE mlx_ext
-    SOURCES ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
-    INCLUDE_DIRS ${PROJECT_SOURCE_DIR} ${MLX_INCLUDE_DIRS}
-    OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}
-  )
-
-  add_dependencies(
-    mlx_ext
+    TARGET
     mlx_ext_metallib
-  )
+    TITLE
+    mlx_ext
+    SOURCES
+    ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
+    INCLUDE_DIRS
+    ${PROJECT_SOURCE_DIR}
+    ${MLX_INCLUDE_DIRS}
+    OUTPUT_DIRECTORY
+    ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
+
+  add_dependencies(mlx_ext mlx_ext_metallib)
 
 endif()
 
 # ----------------------------- Python Bindings -----------------------------
 nanobind_add_module(
   _ext
-  NB_STATIC STABLE_ABI LTO NOMINSIZE
-  NB_DOMAIN mlx 
-  ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
-)
+  NB_STATIC
+  STABLE_ABI
+  LTO
+  NOMINSIZE
+  NB_DOMAIN
+  mlx
+  ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp)
 target_link_libraries(_ext PRIVATE mlx_ext)
 
 if(BUILD_SHARED_LIBS)

examples/extensions/axpby/axpby.cpp

@@ -19,7 +19,7 @@
 #include "mlx/backend/metal/utils.h"
 #endif
 
-namespace mlx::core {
+namespace my_ext {
 
 ///////////////////////////////////////////////////////////////////////////////
 // Operation Implementation
@@ -32,24 +32,24 @@ namespace mlx::core {
  *  Follow numpy style broadcasting between x and y
  *  Inputs are upcasted to floats if needed
  **/
-array axpby(
-    const array& x, // Input array x
-    const array& y, // Input array y
+mx::array axpby(
+    const mx::array& x, // Input mx::array x
+    const mx::array& y, // Input mx::array y
     const float alpha, // Scaling factor for x
     const float beta, // Scaling factor for y
-    StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
+    mx::StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
 ) {
   // Promote dtypes between x and y as needed
   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 = mx::issubdtype(promoted_dtype, mx::float32)
       ? promoted_dtype
-      : promote_types(promoted_dtype, float32);
+      : promote_types(promoted_dtype, mx::float32);
 
   // Cast x and y up to the determined dtype (on the same stream s)
-  auto x_casted = astype(x, out_dtype, s);
-  auto y_casted = astype(y, out_dtype, s);
+  auto x_casted = mx::astype(x, out_dtype, s);
+  auto y_casted = mx::astype(y, out_dtype, s);
 
   // Broadcast the shapes of x and y (on the same stream s)
   auto broadcasted_inputs = broadcast_arrays({x_casted, y_casted}, s);
@@ -57,12 +57,12 @@ array axpby(
 
   // Construct the array as the output of the Axpby primitive
   // with the broadcasted and upcasted arrays as inputs
-  return array(
-      /* const std::vector<int>& shape = */ out_shape,
-      /* Dtype dtype = */ out_dtype,
-      /* std::unique_ptr<Primitive> primitive = */
+  return mx::array(
+      /* const mx::Shape& shape = */ out_shape,
+      /* mx::Dtype dtype = */ out_dtype,
+      /* std::shared_ptr<mx::Primitive> primitive = */
       std::make_shared<Axpby>(to_stream(s), alpha, beta),
-      /* const std::vector<array>& inputs = */ broadcasted_inputs);
+      /* const std::vector<mx::array>& inputs = */ broadcasted_inputs);
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -71,16 +71,16 @@ array axpby(
 
 template <typename T>
 void axpby_impl(
-    const array& x,
-    const array& y,
-    array& out,
+    const mx::array& x,
+    const mx::array& y,
+    mx::array& out,
     float alpha_,
     float beta_) {
   // We only allocate memory when we are ready to fill the output
   // malloc_or_wait synchronously allocates available memory
   // There may be a wait executed here if the allocation is requested
   // under memory-pressured conditions
-  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
 
   // Collect input and output data pointers
   const T* x_ptr = x.data<T>();
@@ -94,8 +94,8 @@ void axpby_impl(
   // Do the element-wise operation for each output
   for (size_t out_idx = 0; out_idx < out.size(); out_idx++) {
     // Map linear indices to offsets in x and y
-    auto x_offset = elem_to_loc(out_idx, x.shape(), x.strides());
-    auto y_offset = elem_to_loc(out_idx, y.shape(), y.strides());
+    auto x_offset = mx::elem_to_loc(out_idx, x.shape(), x.strides());
+    auto y_offset = mx::elem_to_loc(out_idx, y.shape(), y.strides());
 
     // We allocate the output to be contiguous and regularly strided
     // (defaults to row major) and hence it doesn't need additional mapping
@@ -105,8 +105,8 @@ void axpby_impl(
 
 /** Fall back implementation for evaluation on CPU */
 void Axpby::eval(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    const std::vector<mx::array>& inputs,
+    std::vector<mx::array>& outputs) {
   // Check the inputs (registered in the op while constructing the out array)
   assert(inputs.size() == 2);
   auto& x = inputs[0];
@@ -114,14 +114,14 @@ void Axpby::eval(
   auto& out = outputs[0];
 
   // Dispatch to the correct dtype
-  if (out.dtype() == float32) {
+  if (out.dtype() == mx::float32) {
     return axpby_impl<float>(x, y, out, alpha_, beta_);
-  } else if (out.dtype() == float16) {
-    return axpby_impl<float16_t>(x, y, out, alpha_, beta_);
-  } else if (out.dtype() == bfloat16) {
-    return axpby_impl<bfloat16_t>(x, y, out, alpha_, beta_);
-  } else if (out.dtype() == complex64) {
-    return axpby_impl<complex64_t>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == mx::float16) {
+    return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == mx::bfloat16) {
+    return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == mx::complex64) {
+    return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_);
   } else {
     throw std::runtime_error(
         "Axpby is only supported for floating point types.");
@@ -136,9 +136,9 @@ void Axpby::eval(
 
 template <typename T>
 void axpby_impl_accelerate(
-    const array& x,
-    const array& y,
-    array& out,
+    const mx::array& x,
+    const mx::array& y,
+    mx::array& out,
     float alpha_,
     float beta_) {
   // Accelerate library provides catlas_saxpby which does
@@ -150,10 +150,10 @@ 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(mx::allocator::malloc_or_wait(out.nbytes()));
 
   // We then copy over the elements using the contiguous vector specialization
-  copy_inplace(y, out, CopyType::Vector);
+  copy_inplace(y, out, mx::CopyType::Vector);
 
   // Get x and y pointers for catlas_saxpby
   const T* x_ptr = x.data<T>();
@@ -175,15 +175,15 @@ void axpby_impl_accelerate(
 
 /** Evaluate primitive on CPU using accelerate specializations */
 void Axpby::eval_cpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    const std::vector<mx::array>& inputs,
+    std::vector<mx::array>& outputs) {
   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 &&
+  if (out.dtype() == mx::float32 &&
       ((x.flags().row_contiguous && y.flags().row_contiguous) ||
        (x.flags().col_contiguous && y.flags().col_contiguous))) {
     axpby_impl_accelerate<float>(x, y, out, alpha_, beta_);
@@ -198,8 +198,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) {
+    const std::vector<mx::array>& inputs,
+    std::vector<mx::array>& outputs) {
   eval(inputs, outputs);
 }
 
@@ -213,8 +213,8 @@ void Axpby::eval_cpu(
 
 /** Evaluate primitive on GPU */
 void Axpby::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
+    const std::vector<mx::array>& inputs,
+    std::vector<mx::array>& outputs) {
   // Prepare inputs
   assert(inputs.size() == 2);
   auto& x = inputs[0];
@@ -225,7 +225,7 @@ void Axpby::eval_gpu(
   // and each stream carries its device identifiers
   auto& s = stream();
   // We get the needed metal device using the stream
-  auto& d = metal::device(s.device);
+  auto& d = mx::metal::device(s.device);
 
   // Prepare to specialize based on contiguity
   bool contiguous_kernel =
@@ -235,12 +235,12 @@ void Axpby::eval_gpu(
   // Allocate output memory with strides based on specialization
   if (contiguous_kernel) {
     out.set_data(
-        allocator::malloc_or_wait(x.data_size() * out.itemsize()),
+        mx::allocator::malloc_or_wait(x.data_size() * out.itemsize()),
         x.data_size(),
         x.strides(),
         x.flags());
   } else {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    out.set_data(mx::allocator::malloc_or_wait(out.nbytes()));
   }
 
   // Resolve name of kernel (corresponds to axpby.metal)
@@ -257,7 +257,7 @@ void Axpby::eval_gpu(
 
   // Prepare to encode kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
-  compute_encoder->setComputePipelineState(kernel);
+  compute_encoder.set_compute_pipeline_state(kernel);
 
   // Kernel parameters are registered with buffer indices corresponding to
   // those in the kernel declaration at axpby.metal
@@ -272,15 +272,15 @@ void Axpby::eval_gpu(
   compute_encoder.set_output_array(out, 2);
 
   // Encode alpha and beta
-  compute_encoder->setBytes(&alpha_, sizeof(float), 3);
-  compute_encoder->setBytes(&beta_, sizeof(float), 4);
+  compute_encoder.set_bytes(alpha_, 3);
+  compute_encoder.set_bytes(beta_, 4);
 
   // Encode shape, strides and ndim if needed
   if (!contiguous_kernel) {
-    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);
-    compute_encoder->setBytes(&ndim, sizeof(int), 8);
+    compute_encoder.set_vector_bytes(x.shape(), 5);
+    compute_encoder.set_vector_bytes(x.strides(), 6);
+    compute_encoder.set_vector_bytes(y.strides(), 7);
+    compute_encoder.set_bytes(ndim, 8);
   }
 
   // We launch 1 thread for each input and make sure that the number of
@@ -295,15 +295,15 @@ 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.dispatch_threads(grid_dims, group_dims);
 }
 
 #else // Metal is not available
 
 /** Fail evaluation on GPU */
 void Axpby::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& out) {
+    const std::vector<mx::array>& inputs,
+    std::vector<mx::array>& out) {
   throw std::runtime_error("Axpby has no GPU implementation.");
 }
 
@@ -314,9 +314,9 @@ void Axpby::eval_gpu(
 ///////////////////////////////////////////////////////////////////////////////
 
 /** The Jacobian-vector product. */
-std::vector<array> Axpby::jvp(
-    const std::vector<array>& primals,
-    const std::vector<array>& tangents,
+std::vector<mx::array> Axpby::jvp(
+    const std::vector<mx::array>& primals,
+    const std::vector<mx::array>& tangents,
     const std::vector<int>& argnums) {
   // Forward mode diff that pushes along the tangents
   // The jvp transform on the primitive can built with ops
@@ -328,8 +328,8 @@ std::vector<array> Axpby::jvp(
   // scaled by beta
   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())};
+    auto scale_arr = mx::array(scale, tangents[0].dtype());
+    return {mx::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
@@ -339,24 +339,24 @@ std::vector<array> Axpby::jvp(
 }
 
 /** The vector-Jacobian product. */
-std::vector<array> Axpby::vjp(
-    const std::vector<array>& primals,
-    const std::vector<array>& cotangents,
+std::vector<mx::array> Axpby::vjp(
+    const std::vector<mx::array>& primals,
+    const std::vector<mx::array>& cotangents,
     const std::vector<int>& argnums,
-    const std::vector<array>&) {
+    const std::vector<mx::array>&) {
   // Reverse mode diff
-  std::vector<array> vjps;
+  std::vector<mx::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 = mx::array(scale, cotangents[0].dtype());
+    vjps.push_back(mx::multiply(scale_arr, cotangents[0], stream()));
   }
   return vjps;
 }
 
 /** Vectorize primitive along given axis */
-std::pair<std::vector<array>, std::vector<int>> Axpby::vmap(
-    const std::vector<array>& inputs,
+std::pair<std::vector<mx::array>, std::vector<int>> Axpby::vmap(
+    const std::vector<mx::array>& inputs,
     const std::vector<int>& axes) {
   throw std::runtime_error("Axpby has no vmap implementation.");
 }
@@ -367,4 +367,4 @@ bool Axpby::is_equivalent(const Primitive& other) const {
   return alpha_ == r_other.alpha_ && beta_ == r_other.beta_;
 }
 
-} // namespace mlx::core
+} // namespace my_ext

examples/extensions/axpby/axpby.h

@@ -5,7 +5,9 @@
 #include "mlx/ops.h"
 #include "mlx/primitives.h"
 
-namespace mlx::core {
+namespace mx = mlx::core;
+
+namespace my_ext {
 
 ///////////////////////////////////////////////////////////////////////////////
 // Operation
@@ -18,22 +20,22 @@ namespace mlx::core {
  *  Follow numpy style broadcasting between x and y
  *  Inputs are upcasted to floats if needed
  **/
-array axpby(
-    const array& x, // Input array x
-    const array& y, // Input array y
+mx::array axpby(
+    const mx::array& x, // Input array x
+    const mx::array& y, // Input array y
     const float alpha, // Scaling factor for x
     const float beta, // Scaling factor for y
-    StreamOrDevice s = {} // Stream on which to schedule the operation
+    mx::StreamOrDevice s = {} // Stream on which to schedule the operation
 );
 
 ///////////////////////////////////////////////////////////////////////////////
 // Primitive
 ///////////////////////////////////////////////////////////////////////////////
 
-class Axpby : public Primitive {
+class Axpby : public mx::Primitive {
  public:
-  explicit Axpby(Stream stream, float alpha, float beta)
-      : Primitive(stream), alpha_(alpha), beta_(beta) {};
+  explicit Axpby(mx::Stream stream, float alpha, float beta)
+      : mx::Primitive(stream), alpha_(alpha), beta_(beta) {};
 
   /**
    * A primitive must know how to evaluate itself on the CPU/GPU
@@ -42,23 +44,25 @@ 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)
-      override;
-  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
-      override;
+  void eval_cpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
 
   /** The Jacobian-vector product. */
-  std::vector<array> jvp(
-      const std::vector<array>& primals,
-      const std::vector<array>& tangents,
+  std::vector<mx::array> jvp(
+      const std::vector<mx::array>& primals,
+      const std::vector<mx::array>& tangents,
       const std::vector<int>& argnums) override;
 
   /** The vector-Jacobian product. */
-  std::vector<array> vjp(
-      const std::vector<array>& primals,
-      const std::vector<array>& cotangents,
+  std::vector<mx::array> vjp(
+      const std::vector<mx::array>& primals,
+      const std::vector<mx::array>& cotangents,
       const std::vector<int>& argnums,
-      const std::vector<array>& outputs) override;
+      const std::vector<mx::array>& outputs) override;
 
   /**
    * The primitive must know how to vectorize itself across
@@ -66,8 +70,8 @@ class Axpby : public Primitive {
    * representing the vectorized computation and the axis which
    * corresponds to the output vectorized dimension.
    */
-  std::pair<std::vector<array>, std::vector<int>> vmap(
-      const std::vector<array>& inputs,
+  std::pair<std::vector<mx::array>, std::vector<int>> vmap(
+      const std::vector<mx::array>& inputs,
       const std::vector<int>& axes) override;
 
   /** Print the primitive. */
@@ -76,14 +80,16 @@ class Axpby : public Primitive {
   }
 
   /** Equivalence check **/
-  bool is_equivalent(const Primitive& other) const override;
+  bool is_equivalent(const mx::Primitive& other) const override;
 
  private:
   float alpha_;
   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<mx::array>& inputs,
+      std::vector<mx::array>& outputs);
 };
 
-} // namespace mlx::core
+} // namespace my_ext

examples/extensions/axpby/axpby.metal

@@ -2,7 +2,6 @@
 
 #include <metal_stdlib>
 
-#include "mlx/backend/metal/kernels/bf16.h"
 #include "mlx/backend/metal/kernels/utils.h"
 
 template <typename T>
@@ -13,8 +12,8 @@ template <typename T>
     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 int64_t* x_strides [[buffer(6)]],
+    constant const int64_t* y_strides [[buffer(7)]],
     constant const int& ndim [[buffer(8)]],
     uint index [[thread_position_in_grid]]) {
   auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
@@ -35,29 +34,14 @@ template <typename T>
       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)]],                          \
-      uint index [[thread_position_in_grid]]);
+// clang-format off
+#define instantiate_axpby(type_name, type)                             \
+  instantiate_kernel("axpby_general_" #type_name, axpby_general, type) \
+  instantiate_kernel(                                                  \
+          "axpby_contiguous_" #type_name, axpby_contiguous, type)
 
 instantiate_axpby(float32, float);
 instantiate_axpby(float16, half);
 instantiate_axpby(bfloat16, bfloat16_t);
-instantiate_axpby(complex64, complex64_t);
+instantiate_axpby(complex64, complex64_t);
+// clang-format on

examples/extensions/bindings.cpp

@@ -8,14 +8,12 @@
 namespace nb = nanobind;
 using namespace nb::literals;
 
-using namespace mlx::core;
-
 NB_MODULE(_ext, m) {
   m.doc() = "Sample extension for MLX";
 
   m.def(
       "axpby",
-      &axpby,
+      &my_ext::axpby,
       "x"_a,
       "y"_a,
       "alpha"_a,

examples/extensions/pyproject.toml

@@ -1,8 +1,8 @@
 [build-system]
 requires = [
   "setuptools>=42",
-  "cmake>=3.24",
-  "mlx>=0.17.0",
-  "nanobind==2.1.0",
+  "cmake>=3.25",
+  "mlx>=0.18.0",
+  "nanobind==2.4.0",
 ]
 build-backend = "setuptools.build_meta"

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
-cmake>=3.24
-mlx>=0.17.0
-nanobind==2.1.0
+cmake>=3.25
+mlx>=0.21.0
+nanobind==2.2.0

mlx.pc.in

@@ -28,10 +28,19 @@ endif()
 if (@MLX_BUILD_METAL@)
     set(MLX_BUILD_METAL @MLX_BUILD_METAL@)
     set(MLX_CXX_FLAGS ${MLX_CXX_FLAGS} -D_METAL_)
-    set_and_check(MLX_INCLUDE_DIRS 
-        ${MLX_INCLUDE_DIRS} 
+    set(MLX_INCLUDE_DIRS 
+        "${MLX_INCLUDE_DIRS};"
         @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/metal_cpp
     )
+    if(@MLX_METAL_VERSION@ GREATER_EQUAL 310)
+      set(MLX_INCLUDE_DIRS
+        "${MLX_INCLUDE_DIRS};"
+        @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/mlx/backend/metal/kernels/metal_3_1)
+    else()
+      set(MLX_INCLUDE_DIRS
+        "${MLX_INCLUDE_DIRS};"
+        @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/mlx/backend/metal/kernels/metal_3_0)
+    endif()
 endif()
 
 set_target_properties(mlx PROPERTIES
@@ -40,4 +49,4 @@ set_target_properties(mlx PROPERTIES
 )
 
 include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(MLX DEFAULT_MSG MLX_LIBRARY MLX_INCLUDE_DIRS)
+find_package_handle_standard_args(MLX DEFAULT_MSG MLX_LIBRARY MLX_INCLUDE_DIRS)

mlx/CMakeLists.txt

@@ -1,26 +1,35 @@
 target_sources(
   mlx
-  PRIVATE
-  ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
-  ${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
-  ${CMAKE_CURRENT_SOURCE_DIR}/graph_utils.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h
-)
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/export.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
+          ${CMAKE_CURRENT_SOURCE_DIR}/graph_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
 
-if (MLX_BUILD_CPU)
+if(MSVC)
+  # Disable some MSVC warnings to speed up compilation.
+  target_compile_options(mlx PUBLIC /wd4068 /wd4244 /wd4267 /wd4804)
+endif()
+
+if(WIN32)
+  # Export symbols by default to behave like macOS/linux.
+  set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE)
+endif()
+
+if(MLX_BUILD_CPU)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
 else()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_cpu)
@@ -28,17 +37,15 @@ endif()
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
-if (MLX_BUILD_ACCELERATE)
+if(MLX_BUILD_ACCELERATE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/accelerate)
 elseif(MLX_BUILD_CPU)
   target_sources(
     mlx
-    PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/backend/common/default_primitives.cpp
-  )
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/common/default_primitives.cpp)
 endif()
 
-if (MLX_BUILD_METAL)
+if(MLX_BUILD_METAL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
 else()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_metal)

mlx/allocator.cpp

@@ -19,7 +19,7 @@ Buffer malloc(size_t size) {
 }
 
 void free(Buffer buffer) {
-  return allocator().free(buffer);
+  allocator().free(buffer);
 }
 
 Buffer CommonAllocator::malloc(size_t size, bool) {

mlx/array.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <functional>
+#include <unordered_map>
 
 #include "mlx/array.h"
 #include "mlx/ops.h"
@@ -24,13 +25,13 @@ bool retain_graph() {
 } // namespace
 
 array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
-    : array_desc_(std::make_shared<ArrayDesc>(std::vector<int>{}, dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(Shape{}, dtype)) {
   auto cval = static_cast<complex64_t>(val);
   init(&cval);
 }
 
 array::array(
-    std::vector<int> shape,
+    Shape shape,
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
     std::vector<array> inputs)
@@ -41,7 +42,7 @@ array::array(
           std::move(inputs))) {}
 
 std::vector<array> array::make_arrays(
-    std::vector<std::vector<int>> shapes,
+    std::vector<Shape> shapes,
     const std::vector<Dtype>& dtypes,
     const std::shared_ptr<Primitive>& primitive,
     const std::vector<array>& inputs) {
@@ -60,24 +61,20 @@ std::vector<array> array::make_arrays(
 
 array::array(std::initializer_list<float> data)
     : array_desc_(std::make_shared<ArrayDesc>(
-          std::vector<int>{static_cast<int>(data.size())},
+          Shape{static_cast<ShapeElem>(data.size())},
           float32)) {
   init(data.begin());
 }
 
 array::array(std::initializer_list<int> data, Dtype dtype)
     : array_desc_(std::make_shared<ArrayDesc>(
-          std::vector<int>{static_cast<int>(data.size())},
+          Shape{static_cast<ShapeElem>(data.size())},
           dtype)) {
   init(data.begin());
 }
 
 /* Build an array from a shared buffer */
-array::array(
-    allocator::Buffer data,
-    std::vector<int> shape,
-    Dtype dtype,
-    deleter_t deleter)
+array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter)
     : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   set_data(data, deleter);
 }
@@ -95,21 +92,37 @@ void array::detach() {
   array_desc_->primitive = nullptr;
 }
 
-void array::eval() {
-  // Ensure the array is ready to be read
-  if (status() == Status::scheduled) {
+bool array::is_available() const {
+  if (status() == Status::available) {
+    return true;
+  } else if (status() == Status::evaluated && event().is_signaled()) {
+    set_status(Status::available);
+    return true;
+  }
+  return false;
+}
+
+void array::wait() {
+  if (!is_available()) {
     event().wait();
     set_status(Status::available);
-  } else if (status() == Status::unscheduled) {
+  }
+}
+
+void array::eval() {
+  // Ensure the array is ready to be read
+  if (status() == Status::unscheduled) {
     mlx::core::eval({*this});
+  } else {
+    wait();
   }
 }
 
 bool array::is_tracer() const {
-  return array_desc_->is_tracer && in_tracing() || retain_graph();
+  return (array_desc_->is_tracer && in_tracing()) || retain_graph();
 }
 
-void array::set_data(allocator::Buffer buffer, deleter_t d) {
+void array::set_data(allocator::Buffer buffer, Deleter d) {
   array_desc_->data = std::make_shared<Data>(buffer, d);
   array_desc_->data_ptr = buffer.raw_ptr();
   array_desc_->data_size = size();
@@ -122,9 +135,9 @@ void array::set_data(allocator::Buffer buffer, deleter_t d) {
 void array::set_data(
     allocator::Buffer buffer,
     size_t data_size,
-    std::vector<size_t> strides,
+    Strides strides,
     Flags flags,
-    deleter_t d) {
+    Deleter d) {
   array_desc_->data = std::make_shared<Data>(buffer, d);
   array_desc_->data_ptr = buffer.raw_ptr();
   array_desc_->data_size = data_size;
@@ -134,7 +147,7 @@ void array::set_data(
 
 void array::copy_shared_buffer(
     const array& other,
-    const std::vector<size_t>& strides,
+    const Strides& strides,
     Flags flags,
     size_t data_size,
     size_t offset /* = 0 */) {
@@ -153,7 +166,7 @@ void array::copy_shared_buffer(const array& other) {
 
 void array::move_shared_buffer(
     array other,
-    const std::vector<size_t>& strides,
+    const Strides& strides,
     Flags flags,
     size_t data_size,
     size_t offset /* = 0 */) {
@@ -162,8 +175,10 @@ void array::move_shared_buffer(
   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);
+  auto data_ptr = other.array_desc_->data_ptr;
+  other.array_desc_->data_ptr = nullptr;
+  array_desc_->data_ptr =
+      static_cast<void*>(static_cast<char*>(data_ptr) + char_offset);
 }
 
 void array::move_shared_buffer(array other) {
@@ -196,6 +211,8 @@ array::~array() {
     if (do_detach) {
       for (auto& s : siblings()) {
         for (auto& ss : s.siblings()) {
+          // Set to null here to avoid descending into array destructor
+          // for siblings
           ss.array_desc_ = nullptr;
         }
         s.array_desc_->siblings.clear();
@@ -216,13 +233,13 @@ void array::ArrayDesc::init() {
   }
 }
 
-array::ArrayDesc::ArrayDesc(std::vector<int> shape, Dtype dtype)
+array::ArrayDesc::ArrayDesc(Shape shape, Dtype dtype)
     : shape(std::move(shape)), dtype(dtype), status(Status::available) {
   init();
 }
 
 array::ArrayDesc::ArrayDesc(
-    std::vector<int> shape,
+    Shape shape,
     Dtype dtype,
     std::shared_ptr<Primitive> primitive,
     std::vector<array> inputs)
@@ -242,25 +259,58 @@ array::ArrayDesc::~ArrayDesc() {
   // 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.
+  if (inputs.empty()) {
+    return;
+  }
+
   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_));
+  auto append_deletable_inputs = [&for_deletion](ArrayDesc& ad) {
+    std::unordered_map<std::uintptr_t, array> input_map;
+    for (array& a : ad.inputs) {
+      if (a.array_desc_) {
+        input_map.insert({a.id(), a});
+        for (auto& s : a.siblings()) {
+          input_map.insert({s.id(), s});
+        }
+      }
     }
-  }
+    ad.inputs.clear();
+    for (auto& [_, a] : input_map) {
+      bool is_deletable =
+          (a.array_desc_.use_count() <= a.siblings().size() + 1);
+      // An array with siblings is deletable only if all of its siblings
+      // are deletable
+      for (auto& s : a.siblings()) {
+        if (!is_deletable) {
+          break;
+        }
+        int is_input = (input_map.find(s.id()) != input_map.end());
+        is_deletable &=
+            s.array_desc_.use_count() <= a.siblings().size() + is_input;
+      }
+      if (is_deletable) {
+        for_deletion.push_back(std::move(a.array_desc_));
+      }
+    }
+  };
+
+  append_deletable_inputs(*this);
 
   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();
+    append_deletable_inputs(*top);
 
-    for (array& a : top->inputs) {
-      if (a.array_desc_.use_count() == 1) {
-        for_deletion.push_back(std::move(a.array_desc_));
-      }
+    // Clear out possible siblings to break circular references
+    for (auto& s : top->siblings) {
+      // Set to null here to avoid descending into top-level
+      // array destructor for siblings
+      s.array_desc_ = nullptr;
     }
+    top->siblings.clear();
   }
 }
 
@@ -272,7 +322,7 @@ array::ArrayIterator::ArrayIterator(const array& arr, int idx)
 }
 
 array::ArrayIterator::reference array::ArrayIterator::operator*() const {
-  auto start = std::vector<int>(arr.ndim(), 0);
+  auto start = Shape(arr.ndim(), 0);
   auto end = arr.shape();
   auto shape = arr.shape();
   shape.erase(shape.begin());

mlx/array.h

@@ -15,7 +15,11 @@ namespace mlx::core {
 
 // Forward declaration
 class Primitive;
-using deleter_t = std::function<void(allocator::Buffer)>;
+
+using Deleter = std::function<void(allocator::Buffer)>;
+using ShapeElem = int32_t;
+using Shape = std::vector<ShapeElem>;
+using Strides = std::vector<int64_t>;
 
 class array {
   /* An array is really a node in a graph. It contains a shared ArrayDesc
@@ -33,7 +37,7 @@ class array {
   template <typename It>
   array(
       It data,
-      std::vector<int> shape,
+      Shape shape,
       Dtype dtype =
           TypeToDtype<typename std::iterator_traits<It>::value_type>());
 
@@ -49,15 +53,15 @@ class array {
   template <typename T>
   array(
       std::initializer_list<T> data,
-      std::vector<int> shape,
+      Shape shape,
       Dtype dtype = TypeToDtype<T>());
 
   /* Build an array from a buffer */
   array(
       allocator::Buffer data,
-      std::vector<int> shape,
+      Shape shape,
       Dtype dtype,
-      deleter_t deleter = allocator::free);
+      Deleter deleter = allocator::free);
 
   /** Assignment to rvalue does not compile. */
   array& operator=(const array& other) && = delete;
@@ -96,7 +100,7 @@ class array {
   }
 
   /** The shape of the array as a vector of integers. */
-  const std::vector<int>& shape() const {
+  const Shape& shape() const {
     return array_desc_->shape;
   }
 
@@ -105,12 +109,12 @@ class array {
    *
    *  This function supports negative indexing and provides
    *  bounds checking. */
-  int shape(int dim) const {
+  auto shape(int dim) const {
     return shape().at(dim < 0 ? dim + ndim() : dim);
   }
 
   /** The strides of the array. */
-  const std::vector<size_t>& strides() const {
+  const Strides& strides() const {
     return array_desc_->strides;
   }
 
@@ -119,7 +123,7 @@ class array {
    *
    *  This function supports negative indexing and provides
    *  bounds checking. */
-  size_t strides(int dim) const {
+  auto strides(int dim) const {
     return strides().at(dim < 0 ? dim + ndim() : dim);
   }
 
@@ -184,13 +188,13 @@ class array {
    */
 
   array(
-      std::vector<int> shape,
+      Shape shape,
       Dtype dtype,
       std::shared_ptr<Primitive> primitive,
       std::vector<array> inputs);
 
   static std::vector<array> make_arrays(
-      std::vector<std::vector<int>> shapes,
+      std::vector<Shape> shapes,
       const std::vector<Dtype>& dtypes,
       const std::shared_ptr<Primitive>& primitive,
       const std::vector<array>& inputs);
@@ -207,8 +211,8 @@ class array {
 
   struct Data {
     allocator::Buffer buffer;
-    deleter_t d;
-    Data(allocator::Buffer buffer, deleter_t d = allocator::free)
+    Deleter d;
+    Data(allocator::Buffer buffer, Deleter d = allocator::free)
         : buffer(buffer), d(d) {}
     // Not copyable
     Data(const Data& d) = delete;
@@ -344,11 +348,33 @@ class array {
     return static_cast<T*>(array_desc_->data_ptr);
   }
 
-  enum Status { unscheduled, scheduled, available };
+  enum Status {
+    // The ouptut of a computation which has not been scheduled.
+    // For example, the status of `x` in `auto x = a + b`.
+    unscheduled,
 
-  bool is_available() const {
-    return status() == Status::available;
-  }
+    // The ouptut of a computation which has been scheduled but `eval_*` has
+    // not yet been called on the array's primitive. A possible
+    // status of `x` in `auto x = a + b; eval(x);`
+    scheduled,
+
+    // The array's `eval_*` function has been run, but the computation is not
+    // necessarily complete. The array will have memory allocated and if it is
+    // not a tracer then it will be detached from the graph.
+    evaluated,
+
+    // If the array is the output of a computation then the computation
+    // is complete. Constant arrays are always available (e.g. `array({1, 2,
+    // 3})`)
+    available
+  };
+
+  // Check if the array is safe to read.
+  bool is_available() const;
+
+  // Wait on the array to be available. After this `is_available` returns
+  // `true`.
+  void wait();
 
   Status status() const {
     return array_desc_->status;
@@ -375,18 +401,18 @@ class array {
   // Check if the array is a tracer array
   bool is_tracer() const;
 
-  void set_data(allocator::Buffer buffer, deleter_t d = allocator::free);
+  void set_data(allocator::Buffer buffer, Deleter d = allocator::free);
 
   void set_data(
       allocator::Buffer buffer,
       size_t data_size,
-      std::vector<size_t> strides,
+      Strides strides,
       Flags flags,
-      deleter_t d = allocator::free);
+      Deleter d = allocator::free);
 
   void copy_shared_buffer(
       const array& other,
-      const std::vector<size_t>& strides,
+      const Strides& strides,
       Flags flags,
       size_t data_size,
       size_t offset = 0);
@@ -395,7 +421,7 @@ class array {
 
   void move_shared_buffer(
       array other,
-      const std::vector<size_t>& strides,
+      const Strides& strides,
       Flags flags,
       size_t data_size,
       size_t offset = 0);
@@ -414,8 +440,8 @@ class array {
   void init(const It src);
 
   struct ArrayDesc {
-    std::vector<int> shape;
-    std::vector<size_t> strides;
+    Shape shape;
+    Strides strides;
     size_t size;
     Dtype dtype;
     std::shared_ptr<Primitive> primitive;
@@ -449,10 +475,10 @@ class array {
     // The arrays position in the output list
     uint32_t position{0};
 
-    explicit ArrayDesc(std::vector<int> shape, Dtype dtype);
+    explicit ArrayDesc(Shape shape, Dtype dtype);
 
     explicit ArrayDesc(
-        std::vector<int> shape,
+        Shape shape,
         Dtype dtype,
         std::shared_ptr<Primitive> primitive,
         std::vector<array> inputs);
@@ -473,14 +499,14 @@ class array {
 
 template <typename T>
 array::array(T val, Dtype dtype /* = TypeToDtype<T>() */)
-    : array_desc_(std::make_shared<ArrayDesc>(std::vector<int>{}, dtype)) {
+    : array_desc_(std::make_shared<ArrayDesc>(Shape{}, dtype)) {
   init(&val);
 }
 
 template <typename It>
 array::array(
   It data,
-  std::vector<int> shape,
+  Shape shape,
   Dtype dtype /* = TypeToDtype<typename std::iterator_traits<It>::value_type>() */) :
     array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   init(data);
@@ -491,7 +517,7 @@ array::array(
     std::initializer_list<T> data,
     Dtype dtype /* = TypeToDtype<T>() */)
     : array_desc_(std::make_shared<ArrayDesc>(
-          std::vector<int>{static_cast<int>(data.size())},
+          Shape{static_cast<ShapeElem>(data.size())},
           dtype)) {
   init(data.begin());
 }
@@ -499,7 +525,7 @@ array::array(
 template <typename T>
 array::array(
     std::initializer_list<T> data,
-    std::vector<int> shape,
+    Shape shape,
     Dtype dtype /* = TypeToDtype<T>() */)
     : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
   if (data.size() != size()) {

mlx/backend/accelerate/CMakeLists.txt

@@ -1,10 +1,8 @@
 target_sources(
   mlx
-  PRIVATE
-  ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
-)
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp)

mlx/backend/accelerate/primitives.cpp

@@ -43,6 +43,7 @@ DEFAULT(NumberOfElements)
 DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
+DEFAULT(ExpandDims)
 DEFAULT(FFT)
 DEFAULT(Floor)
 DEFAULT(Gather)
@@ -65,7 +66,6 @@ DEFAULT(Pad)
 DEFAULT(Partition)
 DEFAULT_MULTI(QRF)
 DEFAULT(RandomBits)
-DEFAULT(Reshape)
 DEFAULT(Remainder)
 DEFAULT(Round)
 DEFAULT(Scatter)
@@ -76,11 +76,13 @@ DEFAULT(Slice)
 DEFAULT(SliceUpdate)
 DEFAULT_MULTI(Split)
 DEFAULT(Sort)
+DEFAULT(Squeeze)
 DEFAULT(StopGradient)
 DEFAULT_MULTI(SVD)
 DEFAULT(Transpose)
 DEFAULT(Inverse)
 DEFAULT(Cholesky)
+DEFAULT_MULTI(Eigh)
 
 void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);

mlx/backend/accelerate/quantized.cpp

@@ -18,49 +18,61 @@ void _qmm_t_4_64(
     const float* biases,
     int M,
     int N,
-    int K) {
+    int K,
+    int B,
+    bool batched_w) {
   constexpr int bits = 4;
   constexpr int group_size = 64;
   constexpr int bitmask = (1 << bits) - 1;
   constexpr int pack_factor = 32 / bits;
   constexpr int packs_in_group = group_size / pack_factor;
 
-  for (int m = 0; m < M; m++) {
-    const uint32_t* w_local = w;
-    const float* scales_local = scales;
-    const float* biases_local = biases;
+  int w_els = N * K / pack_factor;
+  int g_els = w_els * pack_factor / group_size;
 
-    for (int n = 0; n < N; n++) {
-      const simd_float16* x_local = (simd_float16*)x;
-      simd_float16 sum = 0;
-      for (int k = 0; k < K; k += group_size) {
-        float scale = *scales_local++;
-        float bias = *biases_local++;
+  for (int i = 0; i < B; i++) {
+    for (int m = 0; m < M; m++) {
+      const uint32_t* w_local = w;
+      const float* scales_local = scales;
+      const float* biases_local = biases;
 
-        for (int kw = 0; kw < packs_in_group; kw += 2) {
-          // TODO: vectorize this properly
-          simd_uint16 wi;
-          for (int e = 0; e < 2; e++) {
-            uint32_t wii = *w_local++;
-            for (int p = 0; p < 8; p++) {
-              wi[e * 8 + p] = wii & bitmask;
-              wii >>= bits;
+      for (int n = 0; n < N; n++) {
+        const simd_float16* x_local = (simd_float16*)x;
+        simd_float16 sum = 0;
+        for (int k = 0; k < K; k += group_size) {
+          float scale = *scales_local++;
+          float bias = *biases_local++;
+
+          for (int kw = 0; kw < packs_in_group; kw += 2) {
+            // TODO: vectorize this properly
+            simd_uint16 wi;
+            for (int e = 0; e < 2; e++) {
+              uint32_t wii = *w_local++;
+              for (int p = 0; p < 8; p++) {
+                wi[e * 8 + p] = wii & bitmask;
+                wii >>= bits;
+              }
             }
-          }
-          simd_float16 wf = simd_float(wi);
-          wf *= scale;
-          wf += bias;
+            simd_float16 wf = simd_float(wi);
+            wf *= scale;
+            wf += bias;
 
-          sum += (*x_local) * wf;
-          x_local++;
+            sum += (*x_local) * wf;
+            x_local++;
+          }
         }
+
+        *result = simd_reduce_add(sum);
+        result++;
       }
 
-      *result = simd_reduce_add(sum);
-      result++;
+      x += K;
+    }
+    if (batched_w) {
+      w += w_els;
+      scales += g_els;
+      biases += g_els;
     }
-
-    x += K;
   }
 }
 
@@ -82,8 +94,10 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   if (condition) {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
     int K = x.shape(-1);
-    int M = x.size() / K;
+    int M = x.shape(-2);
     int N = out.shape(-1);
+    int B = x.size() / K / M;
+    bool batched_w = w.ndim() > 2;
     _qmm_t_4_64(
         out.data<float>(),
         x.data<float>(),
@@ -92,7 +106,9 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
         biases.data<float>(),
         M,
         N,
-        K);
+        K,
+        B,
+        batched_w);
   } else {
     eval(inputs, out);
   }

mlx/backend/accelerate/softmax.cpp

@@ -33,8 +33,8 @@ namespace {
  * Note: The implementation below is a general fast exp. There could be faster
  *       implementations for numbers strictly < 0.
  */
-inline simd_float16 simd_fast_exp(simd_float16 x) {
-  x *= 1.442695; // multiply with log_2(e)
+inline simd_float16 simd_fast_exp(simd_float16 x_init) {
+  auto x = x_init * 1.442695; // multiply with log_2(e)
   simd_float16 ipart, fpart;
   simd_int16 epart;
   x = simd_clamp(x, -80, 80);
@@ -53,7 +53,9 @@ inline simd_float16 simd_fast_exp(simd_float16 x) {
   // bitshifting
   epart = (simd_int(ipart) + 127) << 23;
 
-  return (*(simd_float16*)&epart) * x;
+  // Avoid supressing NaNs
+  simd_int16 eq = (x_init == x_init);
+  return simd_bitselect(x_init, (*(simd_float16*)&epart) * x, eq);
 }
 
 #if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC

mlx/backend/common/CMakeLists.txt

@@ -1,78 +1,71 @@
-
-if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
   set(COMPILER ${CMAKE_C_COMPILER})
   set(CLANG TRUE)
 else()
   set(COMPILER ${CMAKE_CXX_COMPILER})
 endif()
 
+if(MSVC)
+  set(SHELL_EXT ps1)
+  set(SHELL_CMD powershell -ExecutionPolicy Bypass -File)
+else()
+  set(SHELL_EXT sh)
+  set(SHELL_CMD /bin/bash)
+endif()
+
 add_custom_command(
-    OUTPUT  compiled_preamble.cpp
-    COMMAND /bin/bash
-              ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
-              ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
-              ${COMPILER}
-              ${PROJECT_SOURCE_DIR}
-              ${CLANG}
+  OUTPUT compiled_preamble.cpp
+  COMMAND
+    ${SHELL_CMD} ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.${SHELL_EXT}
+    ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp ${COMPILER}
+    ${PROJECT_SOURCE_DIR} ${CLANG} ${CMAKE_SYSTEM_PROCESSOR}
+  DEPENDS make_compiled_preamble.${SHELL_EXT}
+          compiled_preamble.h
+          ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
+          ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
+          ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
+          ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
+          ops.h)
 
-    DEPENDS make_compiled_preamble.sh
-            compiled_preamble.h
-            ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
-            ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
-            ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
-            ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
-            ops.h
-)
-
-add_custom_target(
-  cpu_compiled_preamble
-  DEPENDS compiled_preamble.cpp
-)
+add_custom_target(cpu_compiled_preamble DEPENDS compiled_preamble.cpp)
 
 add_dependencies(mlx cpu_compiled_preamble)
 
 target_sources(
   mlx
-  PRIVATE
-  ${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}/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
-)
+  PRIVATE ${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}/eigh.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}/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_SOURCE_DIR}/utils.cpp
+          ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp)
 
-if (IOS)
-  target_sources(
-    mlx
-    PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/compiled_nocpu.cpp
-  )
+if(IOS)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/compiled_nocpu.cpp)
 else()
-  target_sources(
-    mlx
-    PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/compiled_cpu.cpp
-  )
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/compiled_cpu.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/jit_compiler.cpp)
 endif()

mlx/backend/common/arg_reduce.cpp

@@ -13,8 +13,8 @@ template <typename InT, typename OpT>
 void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
   auto axis_size = in.shape()[axis];
   auto axis_stride = in.strides()[axis];
-  std::vector<size_t> strides = in.strides();
-  std::vector<int> shape = in.shape();
+  Strides strides = in.strides();
+  Shape shape = in.shape();
   strides.erase(strides.begin() + axis);
   shape.erase(shape.begin() + axis);
   for (uint32_t i = 0; i < out.size(); ++i) {

mlx/backend/common/binary.h

@@ -28,8 +28,8 @@ BinaryOpType get_binary_op_type(const array& a, const array& b) {
   } else if (b.data_size() == 1 && a.flags().contiguous) {
     bopt = BinaryOpType::VectorScalar;
   } else if (
-      a.flags().row_contiguous && b.flags().row_contiguous ||
-      a.flags().col_contiguous && b.flags().col_contiguous) {
+      (a.flags().row_contiguous && b.flags().row_contiguous) ||
+      (a.flags().col_contiguous && b.flags().col_contiguous)) {
     bopt = BinaryOpType::VectorVector;
   } else {
     bopt = BinaryOpType::General;
@@ -122,19 +122,7 @@ void set_binary_op_output_data(
   }
 }
 
-struct UseDefaultBinaryOp {
-  template <typename T, typename U>
-  void operator()(const T* a, const T* b, U* dst, int size) {
-    // Should we throw? This should normally never be called.
-    assert(false);
-  }
-
-  template <typename T, typename U>
-  void operator()(const T* a, const T* b, U* dst_a, U* dst_b, int size) {
-    // Should we throw? This should normally never be called.
-    assert(false);
-  }
-};
+struct UseDefaultBinaryOp {};
 
 template <typename T, typename U, typename Op>
 struct DefaultVectorScalar {
@@ -150,18 +138,6 @@ struct DefaultVectorScalar {
       a++;
     }
   }
-
-  void operator()(const T* a, const T* b, U* dst_a, U* dst_b, int size) {
-    T scalar = *b;
-    while (size-- > 0) {
-      auto dst = op(*a, scalar);
-      *dst_a = dst.first;
-      *dst_b = dst.second;
-      dst_a++;
-      dst_b++;
-      a++;
-    }
-  }
 };
 
 template <typename T, typename U, typename Op>
@@ -178,18 +154,6 @@ struct DefaultScalarVector {
       b++;
     }
   }
-
-  void operator()(const T* a, const T* b, U* dst_a, U* dst_b, int size) {
-    T scalar = *a;
-    while (size-- > 0) {
-      auto dst = op(scalar, *b);
-      *dst_a = dst.first;
-      *dst_b = dst.second;
-      dst_a++;
-      dst_b++;
-      b++;
-    }
-  }
 };
 
 template <typename T, typename U, typename Op>
@@ -206,204 +170,110 @@ struct DefaultVectorVector {
       b++;
     }
   }
-
-  void operator()(const T* a, const T* b, U* dst_a, U* dst_b, int size) {
-    while (size-- > 0) {
-      auto dst = op(*a, *b);
-      *dst_a = dst.first;
-      *dst_b = dst.second;
-      dst_a++;
-      dst_b++;
-      a++;
-      b++;
-    }
-  }
 };
 
-template <typename T, typename U, typename Op>
-void binary_op_dims1(const array& a, const array& b, array& out, Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < out.size(); ++i) {
-    dst[i] = op(a_ptr[a_idx], b_ptr[b_idx]);
-    a_idx += a.strides()[0];
-    b_idx += b.strides()[0];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims1(
-    const array& a,
-    const array& b,
-    array& out,
+template <typename T, typename U, typename Op, int D, bool Strided>
+void binary_op_dims(
+    const T* a,
+    const T* b,
+    U* out,
     Op op,
-    int stride) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; i++) {
-    op(a_ptr + a_idx, b_ptr + b_idx, dst, stride);
-    a_idx += a.strides()[0];
-    b_idx += b.strides()[0];
-    dst += stride;
-  }
-}
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
 
-template <typename T, typename U, typename Op>
-void binary_op_dims2(const array& a, const array& b, array& out, Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx]);
-      a_idx += a.strides()[1];
-      b_idx += b.strides()[1];
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims2(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op,
-    int stride) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      op(a_ptr + a_idx, b_ptr + b_idx, dst, stride);
-      a_idx += a.strides()[1];
-      b_idx += b.strides()[1];
-      dst += stride;
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims3(const array& a, const array& b, array& out, Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      for (size_t k = 0; k < a.shape()[2]; ++k) {
-        dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx]);
-        a_idx += a.strides()[2];
-        b_idx += b.strides()[2];
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      binary_op_dims<T, U, Op, D - 1, Strided>(
+          a, b, out, op, shape, a_strides, b_strides, out_strides, axis + 1);
+    } else {
+      if constexpr (Strided) {
+        op(a, b, out, stride_out);
+      } else {
+        *out = op(*a, *b);
       }
-      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
-      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
     }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
+    out += stride_out;
+    a += stride_a;
+    b += stride_b;
   }
 }
 
-template <typename T, typename U, typename Op>
-void binary_op_dims4(const array& a, const array& b, array& out, Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      for (size_t k = 0; k < a.shape()[2]; ++k) {
-        for (size_t ii = 0; ii < a.shape()[3]; ++ii) {
-          dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx]);
-          a_idx += a.strides()[3];
-          b_idx += b.strides()[3];
-        }
-        a_idx += a.strides()[2] - a.strides()[3] * a.shape()[3];
-        b_idx += b.strides()[2] - b.strides()[3] * b.shape()[3];
-      }
-      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
-      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dispatch_dims(
-    const array& a,
-    const array& b,
-    array& out,
-    Op op) {
-  switch (out.ndim()) {
-    case 1:
-      binary_op_dims1<T, U, Op>(a, b, out, op);
-      return;
-    case 2:
-      binary_op_dims2<T, U, Op>(a, b, out, op);
-      return;
-    case 3:
-      binary_op_dims3<T, U, Op>(a, b, out, op);
-      return;
-    case 4:
-      binary_op_dims4<T, U, Op>(a, b, out, op);
-      return;
-  }
-
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  for (size_t i = 0; i < out.size(); i++) {
-    int a_idx = elem_to_loc(i, a.shape(), a.strides());
-    int b_idx = elem_to_loc(i, b.shape(), b.strides());
-    dst[i] = op(a_ptr[a_idx], b_ptr[b_idx]);
-  }
-}
-
-template <typename T, typename U, typename Op>
+template <typename T, typename U, bool Strided, typename Op>
 void binary_op_dispatch_dims(
     const array& a,
     const array& b,
     array& out,
     Op op,
     int dim,
-    int stride) {
-  // Number of dimensions to loop over for vectorized ops
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides) {
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  U* out_ptr = out.data<U>();
   switch (dim) {
     case 1:
-      binary_op_dims1<T, U, Op>(a, b, out, op, stride);
+      binary_op_dims<T, U, Op, 1, Strided>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
       return;
     case 2:
-      binary_op_dims2<T, U, Op>(a, b, out, op, stride);
+      binary_op_dims<T, U, Op, 2, Strided>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
+      return;
+    case 3:
+      binary_op_dims<T, U, Op, 3, Strided>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
       return;
   }
 
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst = out.data<U>();
-  for (size_t i = 0; i < out.size(); i += stride) {
-    int a_idx = elem_to_loc(i, a.shape(), a.strides());
-    int b_idx = elem_to_loc(i, b.shape(), b.strides());
-    op(a_ptr + a_idx, b_ptr + b_idx, dst, stride);
-    dst += stride;
+  ContiguousIterator a_it(shape, a_strides, dim - 3);
+  ContiguousIterator b_it(shape, b_strides, dim - 3);
+  auto stride = out_strides[dim - 4];
+  for (int64_t elem = 0; elem < a.size(); elem += stride) {
+    binary_op_dims<T, U, Op, 3, Strided>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        out_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        out_strides,
+        dim - 3);
+    a_it.step();
+    b_it.step();
   }
 }
 
@@ -450,29 +320,33 @@ void binary_op(
   }
 
   // General computation so let's try to optimize
+  auto [new_shape, new_strides] = collapse_contiguous_dims(
+      a.shape(), {a.strides(), b.strides(), out.strides()});
+  const auto& a_strides = new_strides[0];
+  const auto& b_strides = new_strides[1];
+  const auto& strides = new_strides[2];
 
   // Get the left-most dim such that the array is row contiguous after
-  auto& strides = out.strides();
-  auto leftmost_rc_dim = [&strides](const array& arr) {
-    int d = arr.ndim() - 1;
-    for (; d >= 0 && arr.strides()[d] == strides[d]; d--) {
+  auto leftmost_rc_dim = [&strides](const auto& arr_strides) {
+    int d = arr_strides.size() - 1;
+    for (; d >= 0 && arr_strides[d] == strides[d]; d--) {
     }
     return d + 1;
   };
-  auto a_rc_dim = leftmost_rc_dim(a);
-  auto b_rc_dim = leftmost_rc_dim(b);
+  auto a_rc_dim = leftmost_rc_dim(a_strides);
+  auto b_rc_dim = leftmost_rc_dim(b_strides);
 
   // Get the left-most dim such that the array is a broadcasted "scalar" after
-  auto leftmost_s_dim = [](const array& arr) {
-    int d = arr.ndim() - 1;
-    for (; d >= 0 && arr.strides()[d] == 0; d--) {
+  auto leftmost_s_dim = [](const auto& arr_strides) {
+    int d = arr_strides.size() - 1;
+    for (; d >= 0 && arr_strides[d] == 0; d--) {
     }
     return d + 1;
   };
-  auto a_s_dim = leftmost_s_dim(a);
-  auto b_s_dim = leftmost_s_dim(b);
+  auto a_s_dim = leftmost_s_dim(a_strides);
+  auto b_s_dim = leftmost_s_dim(b_strides);
 
-  auto ndim = out.ndim();
+  auto ndim = new_shape.size();
 
   // Case 1: LxM and FxM where L and F are broadcastable and M is row contiguous
   int dim = ndim;
@@ -494,27 +368,27 @@ void binary_op(
   // Can be sure dim > 0 since otherwise we would have used one of the fully
   // contiguous methods above. Except for the case that the flags do not
   // correspond to the underlying contiguity.
-  size_t stride;
   if (dim == 0 || strides[dim - 1] < 16) {
-    stride = 1;
     bopt = BinaryOpType::General;
     dim = ndim;
-  } else {
-    stride = strides[dim - 1];
   }
 
   switch (bopt) {
     case BinaryOpType::VectorVector:
-      binary_op_dispatch_dims<T, U>(a, b, out, opvv, dim, stride);
+      binary_op_dispatch_dims<T, U, true>(
+          a, b, out, opvv, dim, new_shape, a_strides, b_strides, strides);
       break;
     case BinaryOpType::VectorScalar:
-      binary_op_dispatch_dims<T, U>(a, b, out, opvs, dim, stride);
+      binary_op_dispatch_dims<T, U, true>(
+          a, b, out, opvs, dim, new_shape, a_strides, b_strides, strides);
       break;
     case BinaryOpType::ScalarVector:
-      binary_op_dispatch_dims<T, U>(a, b, out, opsv, dim, stride);
+      binary_op_dispatch_dims<T, U, true>(
+          a, b, out, opsv, dim, new_shape, a_strides, b_strides, strides);
       break;
     default:
-      binary_op_dispatch_dims<T, U>(a, b, out, op);
+      binary_op_dispatch_dims<T, U, false>(
+          a, b, out, op, dim, new_shape, a_strides, b_strides, strides);
       break;
   }
 }
@@ -531,9 +405,9 @@ void binary_op(
   // TODO: The following mess of constexpr evaluations can probably be achieved
   //       with template specializations and overloading. Would it be simpler?
 
-  if (std::is_same<decltype(opsv), UseDefaultBinaryOp>::value) {
-    if (std::is_same<decltype(opvs), UseDefaultBinaryOp>::value) {
-      if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
+  if constexpr (std::is_same<decltype(opsv), UseDefaultBinaryOp>::value) {
+    if constexpr (std::is_same<decltype(opvs), UseDefaultBinaryOp>::value) {
+      if constexpr (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
         // All ops are UseDefaultBinaryOp (why oh why would someone call that?)
         binary_op<T, T>(
             a,
@@ -554,7 +428,8 @@ void binary_op(
             DefaultVectorScalar<T, T, Op>(op),
             opvv);
       }
-    } else if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
+    } else if constexpr (std::is_same<decltype(opvv), UseDefaultBinaryOp>::
+                             value) {
       // opsv and opvv were UseDefaultBinaryOp
       binary_op<T, T>(
           a,
@@ -569,7 +444,8 @@ void binary_op(
       binary_op<T, T>(
           a, b, out, op, DefaultScalarVector<T, T, Op>(op), opvs, opvv);
     }
-  } else if (std::is_same<decltype(opvs), UseDefaultBinaryOp>::value) {
+  } else if constexpr (std::is_same<decltype(opvs), UseDefaultBinaryOp>::
+                           value) {
     if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
       // opvs and opvv were UseDefaultBinaryOp
       binary_op<T, T>(
@@ -585,7 +461,8 @@ void binary_op(
       binary_op<T, T>(
           a, b, out, op, opsv, DefaultVectorScalar<T, T, Op>(op), opvv);
     }
-  } else if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
+  } else if constexpr (std::is_same<decltype(opvv), UseDefaultBinaryOp>::
+                           value) {
     // opvv was UseDefaultBinaryOp
     binary_op<T, T>(
         a, b, out, op, opsv, opvs, DefaultVectorVector<T, T, Op>(op));

mlx/backend/common/binary_two.h

@@ -9,168 +9,43 @@ namespace mlx::core {
 
 namespace {
 
-template <typename T, typename U, typename Op>
-void binary_op_dims1(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < out_a.size(); ++i) {
-    auto dst = op(a_ptr[a_idx], b_ptr[b_idx]);
-    dst_a[i] = dst.first;
-    dst_b[i] = dst.second;
-    a_idx += a.strides()[0];
-    b_idx += b.strides()[0];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims1(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
+template <typename T, typename U, typename Op, int D>
+void binary_op_dims(
+    const T* a,
+    const T* b,
+    U* out_a,
+    U* out_b,
     Op op,
-    int stride) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; i++) {
-    op(a_ptr + a_idx, b_ptr + b_idx, dst_a, dst_b, stride);
-    a_idx += a.strides()[0];
-    b_idx += b.strides()[0];
-    dst_a += stride;
-    dst_b += stride;
-  }
-}
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
 
-template <typename T, typename U, typename Op>
-void binary_op_dims2(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      auto dst = op(a_ptr[a_idx], b_ptr[b_idx]);
-      dst_a[out_idx] = dst.first;
-      dst_b[out_idx++] = dst.second;
-      a_idx += a.strides()[1];
-      b_idx += b.strides()[1];
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      binary_op_dims<T, U, Op, D - 1>(
+          a,
+          b,
+          out_a,
+          out_b,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          axis + 1);
+    } else {
+      std::tie(*out_a, *out_b) = op(*a, *b);
     }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims2(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op,
-    int stride) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      op(a_ptr + a_idx, b_ptr + b_idx, dst_a, dst_b, stride);
-      a_idx += a.strides()[1];
-      b_idx += b.strides()[1];
-      dst_a += stride;
-      dst_b += stride;
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims3(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      for (size_t k = 0; k < a.shape()[2]; ++k) {
-        auto dst = op(a_ptr[a_idx], b_ptr[b_idx]);
-        dst_a[out_idx] = dst.first;
-        dst_b[out_idx++] = dst.second;
-        a_idx += a.strides()[2];
-        b_idx += b.strides()[2];
-      }
-      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
-      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
-  }
-}
-
-template <typename T, typename U, typename Op>
-void binary_op_dims4(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op) {
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  size_t a_idx = 0;
-  size_t b_idx = 0;
-  size_t out_idx = 0;
-  for (size_t i = 0; i < a.shape()[0]; ++i) {
-    for (size_t j = 0; j < a.shape()[1]; ++j) {
-      for (size_t k = 0; k < a.shape()[2]; ++k) {
-        for (size_t ii = 0; ii < a.shape()[3]; ++ii) {
-          auto dst = op(a_ptr[a_idx], b_ptr[b_idx]);
-          dst_a[out_idx] = dst.first;
-          dst_b[out_idx++] = dst.second;
-          a_idx += a.strides()[3];
-          b_idx += b.strides()[3];
-        }
-        a_idx += a.strides()[2] - a.strides()[3] * a.shape()[3];
-        b_idx += b.strides()[2] - b.strides()[3] * b.shape()[3];
-      }
-      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
-      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
-    }
-    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
-    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
+    a += stride_a;
+    b += stride_b;
+    out_a += stride_out;
+    out_b += stride_out;
   }
 }
 
@@ -181,352 +56,160 @@ void binary_op_dispatch_dims(
     array& out_a,
     array& out_b,
     Op op) {
-  switch (out_a.ndim()) {
+  auto [shape, strides] = collapse_contiguous_dims(
+      a.shape(), {a.strides(), b.strides(), out_a.strides()});
+  const auto& a_strides = strides[0];
+  const auto& b_strides = strides[1];
+  const auto& out_strides = strides[2];
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  U* out_a_ptr = out_a.data<U>();
+  U* out_b_ptr = out_b.data<U>();
+
+  int ndim = shape.size();
+  switch (ndim) {
     case 1:
-      binary_op_dims1<T, U, Op>(a, b, out_a, out_b, op);
+      binary_op_dims<T, U, Op, 1>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
       return;
     case 2:
-      binary_op_dims2<T, U, Op>(a, b, out_a, out_b, op);
-      return;
-    case 3:
-      binary_op_dims3<T, U, Op>(a, b, out_a, out_b, op);
-      return;
-    case 4:
-      binary_op_dims4<T, U, Op>(a, b, out_a, out_b, op);
+      binary_op_dims<T, U, Op, 2>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
       return;
   }
 
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  for (size_t i = 0; i < out_a.size(); i++) {
-    int a_idx = elem_to_loc(i, a.shape(), a.strides());
-    int b_idx = elem_to_loc(i, b.shape(), b.strides());
-    std::tie(dst_a[i], dst_b[i]) = op(a_ptr[a_idx], b_ptr[b_idx]);
+  ContiguousIterator a_it(shape, a_strides, ndim - 2);
+  ContiguousIterator b_it(shape, b_strides, ndim - 2);
+  auto stride = out_strides[ndim - 3];
+  for (size_t elem = 0; elem < a.size(); elem += stride) {
+    binary_op_dims<T, U, Op, 2>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        out_a_ptr + elem,
+        out_b_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        out_strides,
+        ndim - 2);
+    a_it.step();
+    b_it.step();
   }
 }
 
-template <typename T, typename U, typename Op>
-void binary_op_dispatch_dims(
-    const array& a,
-    const array& b,
-    array& out_a,
-    array& out_b,
-    Op op,
-    int dim,
-    int stride) {
-  // Number of dimensions to loop over for vectorized ops
-  switch (dim) {
-    case 1:
-      binary_op_dims1<T, U, Op>(a, b, out_a, out_b, op, stride);
-      return;
-    case 2:
-      binary_op_dims2<T, U, Op>(a, b, out_a, out_b, op, stride);
-      return;
-  }
-
-  const T* a_ptr = a.data<T>();
-  const T* b_ptr = b.data<T>();
-  U* dst_a = out_a.data<U>();
-  U* dst_b = out_b.data<U>();
-  for (size_t i = 0; i < out_a.size(); i += stride) {
-    int a_idx = elem_to_loc(i, a.shape(), a.strides());
-    int b_idx = elem_to_loc(i, b.shape(), b.strides());
-    op(a_ptr + a_idx, b_ptr + b_idx, dst_a, dst_b, stride);
-    dst_a += stride;
-    dst_b += stride;
-  }
-}
-
-template <
-    typename T,
-    typename U,
-    typename Op,
-    typename OpSV,
-    typename OpVS,
-    typename OpVV>
+template <typename T, typename U = T, typename Op>
 void binary_op(
     const array& a,
     const array& b,
-    array& out_a,
-    array& out_b,
-    Op op,
-    OpSV opsv,
-    OpVS opvs,
-    OpVV opvv) {
+    std::vector<array>& outputs,
+    Op op) {
   auto bopt = get_binary_op_type(a, b);
+  auto& out_a = outputs[0];
+  auto& out_b = outputs[1];
   set_binary_op_output_data(a, b, out_a, bopt);
   set_binary_op_output_data(a, b, out_b, bopt);
 
   // The full computation is scalar scalar so call the base op once
+  if (bopt == BinaryOpType::General) {
+    binary_op_dispatch_dims<T, U, Op>(a, b, out_a, out_b, op);
+    return;
+  }
+
+  auto a_ptr = a.data<T>();
+  auto b_ptr = b.data<T>();
+  auto out_a_ptr = out_a.data<U>();
+  auto out_b_ptr = out_b.data<U>();
   if (bopt == BinaryOpType::ScalarScalar) {
-    std::tie(*(out_a.data<U>()), *(out_b.data<U>())) =
-        op(*a.data<T>(), *b.data<T>());
-    return;
-  }
-
-  // The full computation is scalar vector so delegate to the op
-  if (bopt == BinaryOpType::ScalarVector) {
-    opsv(
-        a.data<T>(),
-        b.data<T>(),
-        out_a.data<U>(),
-        out_b.data<U>(),
-        b.data_size());
-    return;
-  }
-
-  // The full computation is vector scalar so delegate to the op
-  if (bopt == BinaryOpType::VectorScalar) {
-    opvs(
-        a.data<T>(),
-        b.data<T>(),
-        out_a.data<U>(),
-        out_b.data<U>(),
-        a.data_size());
-    return;
-  }
-
-  // The full computation is vector vector so delegate to the op
-  if (bopt == BinaryOpType::VectorVector) {
-    opvv(
-        a.data<T>(),
-        b.data<T>(),
-        out_a.data<U>(),
-        out_b.data<U>(),
-        out_a.size());
-    return;
-  }
-
-  // General computation so let's try to optimize
-
-  // Get the left-most dim such that the array is row contiguous after
-  auto& strides = out_a.strides();
-  auto leftmost_rc_dim = [&strides](const array& arr) {
-    int d = arr.ndim() - 1;
-    for (; d >= 0 && arr.strides()[d] == strides[d]; d--) {
+    std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+  } else if (bopt == BinaryOpType::ScalarVector) {
+    for (size_t i = 0; i < b.size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      b_ptr++;
     }
-    return d + 1;
-  };
-  auto a_rc_dim = leftmost_rc_dim(a);
-  auto b_rc_dim = leftmost_rc_dim(b);
-
-  // Get the left-most dim such that the array is a broadcasted "scalar" after
-  auto leftmost_s_dim = [](const array& arr) {
-    int d = arr.ndim() - 1;
-    for (; d >= 0 && arr.strides()[d] == 0; d--) {
+  } else if (bopt == BinaryOpType::VectorScalar) {
+    for (size_t i = 0; i < a.size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+    }
+  } else { // VectorVector
+    for (size_t i = 0; i < a.size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+      b_ptr++;
     }
-    return d + 1;
-  };
-  auto a_s_dim = leftmost_s_dim(a);
-  auto b_s_dim = leftmost_s_dim(b);
-
-  auto ndim = out_a.ndim();
-
-  // Case 1: LxM and FxM where L and F are broadcastable and M is row contiguous
-  int dim = ndim;
-  if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
-    bopt = BinaryOpType::VectorVector;
-    dim = d;
-    // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
-    // contiguous
-  } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
-    bopt = BinaryOpType::VectorScalar;
-    dim = d;
-    // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
-    // contiguous
-  } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
-    bopt = BinaryOpType::ScalarVector;
-    dim = d;
-  }
-
-  // Can be sure dim > 0 since otherwise we would have used one of the fully
-  // contiguous methods above. Except for the case that the flags do not
-  // correspond to the underlying contiguity.
-  size_t stride;
-  if (dim == 0 || strides[dim - 1] < 16) {
-    stride = 1;
-    bopt = BinaryOpType::General;
-    dim = ndim;
-  } else {
-    stride = strides[dim - 1];
-  }
-
-  switch (bopt) {
-    case BinaryOpType::VectorVector:
-      binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opvv, dim, stride);
-      break;
-    case BinaryOpType::VectorScalar:
-      binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opvs, dim, stride);
-      break;
-    case BinaryOpType::ScalarVector:
-      binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opsv, dim, stride);
-      break;
-    default:
-      binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, op);
-      break;
   }
 }
 
-template <typename T, typename Op, typename OpSV, typename OpVS, typename OpVV>
-void binary_op(
-    const array& a,
-    const array& b,
-    std::vector<array>& outputs,
-    Op op,
-    OpSV opsv,
-    OpVS opvs,
-    OpVV opvv) {
-  // TODO: The following mess of constexpr evaluations can probably be achieved
-  //       with template specializations and overloading. Would it be simpler?
-
-  if (std::is_same<decltype(opsv), UseDefaultBinaryOp>::value) {
-    if (std::is_same<decltype(opvs), UseDefaultBinaryOp>::value) {
-      if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
-        // All ops are UseDefaultBinaryOp (why oh why would someone call that?)
-        binary_op<T, T>(
-            a,
-            b,
-            outputs[0],
-            outputs[1],
-            op,
-            DefaultScalarVector<T, T, Op>(op),
-            DefaultVectorScalar<T, T, Op>(op),
-            DefaultVectorVector<T, T, Op>(op));
-      } else {
-        // opsv and opvs were UseDefaultBinaryOp
-        binary_op<T, T>(
-            a,
-            b,
-            outputs[0],
-            outputs[1],
-            op,
-            DefaultScalarVector<T, T, Op>(op),
-            DefaultVectorScalar<T, T, Op>(op),
-            opvv);
-      }
-    } else if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
-      // opsv and opvv were UseDefaultBinaryOp
-      binary_op<T, T>(
-          a,
-          b,
-          outputs[0],
-          outputs[1],
-          op,
-          DefaultScalarVector<T, T, Op>(op),
-          opvs,
-          DefaultVectorVector<T, T, Op>(op));
-    } else {
-      // opsv was UseDefaultBinaryOp
-      binary_op<T, T>(
-          a,
-          b,
-          outputs[0],
-          outputs[1],
-          op,
-          DefaultScalarVector<T, T, Op>(op),
-          opvs,
-          opvv);
-    }
-  } else if (std::is_same<decltype(opvs), UseDefaultBinaryOp>::value) {
-    if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
-      // opvs and opvv were UseDefaultBinaryOp
-      binary_op<T, T>(
-          a,
-          b,
-          outputs[0],
-          outputs[1],
-          op,
-          opsv,
-          DefaultVectorScalar<T, T, Op>(op),
-          DefaultVectorVector<T, T, Op>(op));
-    } else {
-      // opvs was UseDefaultBinaryOp
-      binary_op<T, T>(
-          a,
-          b,
-          outputs[0],
-          outputs[1],
-          op,
-          opsv,
-          DefaultVectorScalar<T, T, Op>(op),
-          opvv);
-    }
-  } else if (std::is_same<decltype(opvv), UseDefaultBinaryOp>::value) {
-    // opvv was UseDefaultBinaryOp
-    binary_op<T, T>(
-        a,
-        b,
-        outputs[0],
-        outputs[1],
-        op,
-        opsv,
-        opvs,
-        DefaultVectorVector<T, T, Op>(op));
-  } else {
-    // All ops provided
-    binary_op<T, T>(a, b, outputs[0], outputs[1], op, opsv, opvs, opvv);
-  }
-}
-
-template <typename T, typename Op>
-void binary_op(
-    const array& a,
-    const array& b,
-    std::vector<array>& outputs,
-    Op op) {
-  DefaultScalarVector<T, T, Op> opsv(op);
-  DefaultVectorScalar<T, T, Op> opvs(op);
-  DefaultVectorVector<T, T, Op> opvv(op);
-  binary_op<T, T>(a, b, outputs[0], outputs[1], op, opsv, opvs, opvv);
-}
-
-template <typename... Ops>
+template <typename Op>
 void binary(
     const array& a,
     const array& b,
     std::vector<array>& outputs,
-    Ops... ops) {
+    Op op) {
   switch (outputs[0].dtype()) {
     case bool_:
-      binary_op<bool>(a, b, outputs, ops...);
+      binary_op<bool>(a, b, outputs, op);
       break;
     case uint8:
-      binary_op<uint8_t>(a, b, outputs, ops...);
+      binary_op<uint8_t>(a, b, outputs, op);
       break;
     case uint16:
-      binary_op<uint16_t>(a, b, outputs, ops...);
+      binary_op<uint16_t>(a, b, outputs, op);
       break;
     case uint32:
-      binary_op<uint32_t>(a, b, outputs, ops...);
+      binary_op<uint32_t>(a, b, outputs, op);
       break;
     case uint64:
-      binary_op<uint64_t>(a, b, outputs, ops...);
+      binary_op<uint64_t>(a, b, outputs, op);
       break;
     case int8:
-      binary_op<int8_t>(a, b, outputs, ops...);
+      binary_op<int8_t>(a, b, outputs, op);
       break;
     case int16:
-      binary_op<int16_t>(a, b, outputs, ops...);
+      binary_op<int16_t>(a, b, outputs, op);
       break;
     case int32:
-      binary_op<int32_t>(a, b, outputs, ops...);
+      binary_op<int32_t>(a, b, outputs, op);
       break;
     case int64:
-      binary_op<int64_t>(a, b, outputs, ops...);
+      binary_op<int64_t>(a, b, outputs, op);
       break;
     case float16:
-      binary_op<float16_t>(a, b, outputs, ops...);
+      binary_op<float16_t>(a, b, outputs, op);
       break;
     case float32:
-      binary_op<float>(a, b, outputs, ops...);
+      binary_op<float>(a, b, outputs, op);
       break;
     case bfloat16:
-      binary_op<bfloat16_t>(a, b, outputs, ops...);
+      binary_op<bfloat16_t>(a, b, outputs, op);
       break;
     case complex64:
-      binary_op<complex64_t>(a, b, outputs, ops...);
+      binary_op<complex64_t>(a, b, outputs, op);
       break;
   }
 }

mlx/backend/common/cholesky.cpp

@@ -2,46 +2,12 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/lapack.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:
@@ -66,7 +32,14 @@ void cholesky_impl(const array& a, array& factor, bool upper) {
 
   for (int i = 0; i < num_matrices; i++) {
     // Compute Cholesky factorization.
-    int info = spotrf_wrapper(uplo, matrix, N);
+    int info;
+    MLX_LAPACK_FUNC(spotrf)
+    (
+        /* uplo = */ &uplo,
+        /* n = */ &N,
+        /* a = */ matrix,
+        /* lda = */ &N,
+        /* info = */ &info);
 
     // 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

mlx/backend/common/common.cpp

@@ -39,7 +39,7 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
   // rely on data_size anyway.
   size_t data_size = out.size();
 
-  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
+  return move_or_copy(in, out, strides_, flags, data_size, offset_);
 }
 
 void Broadcast::eval(const std::vector<array>& inputs, array& out) {
@@ -49,7 +49,7 @@ void Broadcast::eval(const std::vector<array>& inputs, array& out) {
     out.set_data(nullptr);
     return;
   }
-  std::vector<size_t> strides(out.ndim(), 0);
+  Strides 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];
@@ -58,12 +58,12 @@ void Broadcast::eval(const std::vector<array>& inputs, array& out) {
   if (out.size() > in.size()) {
     flags.row_contiguous = flags.col_contiguous = false;
   }
-  out.copy_shared_buffer(in, strides, flags, in.data_size());
+  move_or_copy(in, out, 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]);
+  move_or_copy(inputs[0], out);
 }
 
 void CustomTransforms::eval(
@@ -72,7 +72,7 @@ void CustomTransforms::eval(
   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]);
+    move_or_copy(inputs[j], outputs[i]);
   }
 }
 
@@ -81,10 +81,20 @@ void Depends::eval(
     std::vector<array>& outputs) {
   assert(inputs.size() > outputs.size());
   for (int i = 0; i < outputs.size(); i++) {
-    outputs[i].copy_shared_buffer(inputs[i]);
+    move_or_copy(inputs[i], outputs[i]);
   }
 }
 
+void ExpandDims::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  auto strides = in.strides();
+  for (auto ax : axes_) {
+    strides.insert(strides.begin() + ax, 1);
+  }
+  move_or_copy(in, out, strides, in.flags(), in.data_size());
+}
+
 void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
@@ -141,9 +151,7 @@ void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
-std::pair<bool, std::vector<size_t>> Reshape::prepare_reshape(
-    const array& in,
-    const array& out) {
+std::pair<bool, Strides> 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()};
@@ -151,17 +159,15 @@ std::pair<bool, std::vector<size_t>> Reshape::prepare_reshape(
 
   // Special case for scalars
   if (in.ndim() == 0) {
-    std::vector<size_t> out_strides(out.ndim(), 0);
-    return {false, out_strides};
+    return {false, Strides(out.ndim(), 0)};
   }
 
   // Firstly let's collapse all the contiguous dimensions of the input
-  auto [shape, _strides] = collapse_contiguous_dims(in);
-  auto& strides = _strides[0];
+  auto [shape, strides] = collapse_contiguous_dims(in);
 
   // If shapes fit exactly in the contiguous dims then no copy is necessary so
   // let's check.
-  std::vector<size_t> out_strides;
+  Strides out_strides;
   bool copy_necessary = false;
   int j = 0;
   for (int i = 0; i < out.ndim(); i++) {
@@ -182,9 +188,9 @@ std::pair<bool, std::vector<size_t>> Reshape::prepare_reshape(
   return {copy_necessary, out_strides};
 }
 
-void Reshape::shared_buffer_reshape(
+void shared_buffer_reshape(
     const array& in,
-    const std::vector<size_t>& out_strides,
+    const Strides& out_strides,
     array& out) {
   auto flags = in.flags();
   if (flags.row_contiguous) {
@@ -195,7 +201,7 @@ void Reshape::shared_buffer_reshape(
     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());
+  move_or_copy(in, out, out_strides, flags, in.data_size());
 }
 
 void Split::eval(
@@ -250,26 +256,28 @@ void Split::eval(
   }
 }
 
-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];
+void Squeeze::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  Strides strides;
+  for (int i = 0, j = 0; i < in.ndim(); ++i) {
+    if (j < axes_.size() && i == axes_[j]) {
+      j++;
+    } else {
+      strides.push_back(in.strides(i));
+    }
   }
-
-  return std::make_tuple(data_offset, inp_strides);
+  move_or_copy(in, out, strides, in.flags(), in.data_size());
 }
 
 void StopGradient::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  out.copy_shared_buffer(inputs[0]);
+  move_or_copy(inputs[0], out);
 }
 
 void Transpose::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  std::vector<size_t> out_strides(out.ndim());
+  Strides out_strides(out.ndim());
   auto& in = inputs[0];
   for (int ax = 0; ax < axes_.size(); ++ax) {
     out_strides[ax] = in.strides()[axes_[ax]];
@@ -286,8 +294,8 @@ void Transpose::eval(const std::vector<array>& inputs, array& out) {
   //   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;
+    int64_t f_stride = 1;
+    int64_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) {
@@ -298,7 +306,7 @@ void Transpose::eval(const std::vector<array>& inputs, array& out) {
       b_stride *= out.shape(ri);
     }
   }
-  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
+  move_or_copy(in, out, out_strides, flags, in.data_size());
 }
 
 } // namespace mlx::core

mlx/backend/common/compiled.cpp

@@ -130,7 +130,7 @@ std::string build_lib_name(
 
 bool compiled_check_contiguity(
     const std::vector<array>& inputs,
-    const std::vector<int>& shape) {
+    const Shape& shape) {
   bool contiguous = true;
   bool all_contig = true;
   bool all_row_contig = true;
@@ -165,7 +165,7 @@ void compiled_allocate_outputs(
     bool move_buffers /* = false */) {
   if (contiguous) {
     int o = 0;
-    std::vector<size_t> strides;
+    Strides strides;
     size_t data_size;
     array::Flags flags;
     for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {

mlx/backend/common/compiled.h

@@ -56,7 +56,7 @@ inline bool is_scalar(const array& x) {
 // 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);
+    const Shape& shape);
 
 // Allocate space for the outputs possibly with input donation
 void compiled_allocate_outputs(

mlx/backend/common/compiled_cpu.cpp

@@ -4,30 +4,18 @@
 #include <filesystem>
 #include <fstream>
 #include <list>
+#include <mutex>
+#include <shared_mutex>
 
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/common/compiled_preamble.h"
+#include "mlx/backend/common/jit_compiler.h"
 #include "mlx/device.h"
 #include "mlx/graph_utils.h"
 
 namespace mlx::core {
 
-// GPU compile is always available if the GPU is available and since we are in
-// this file CPU compile is also available.
-namespace detail {
-bool compile_available_for_device(const Device& device) {
-  return true;
-}
-} // namespace detail
-
-std::string get_temp_file(const std::string& name) {
-  return std::filesystem::temp_directory_path().append(name);
-}
-
-// Return a pointer to a compiled function
-void* compile(
-    const std::string& kernel_name,
-    const std::string& source_code = "") {
+struct CompilerCache {
   struct DLib {
     DLib(const std::string& libname) {
       lib = dlopen(libname.c_str(), RTLD_NOW);
@@ -44,35 +32,64 @@ void* compile(
     void* lib;
   };
   // Statics to cache compiled libraries and functions
-  static std::list<DLib> libs;
-  static std::unordered_map<std::string, void*> kernels;
-  if (auto it = kernels.find(kernel_name); it != kernels.end()) {
+  std::list<DLib> libs;
+  std::unordered_map<std::string, void*> kernels;
+  std::shared_mutex mtx;
+};
+
+static CompilerCache cache{};
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is also available.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return true;
+}
+
+} // namespace detail
+
+// Return a pointer to a compiled function
+void* compile(
+    const std::string& kernel_name,
+    const std::function<std::string(void)>& source_builder) {
+  {
+    std::shared_lock lock(cache.mtx);
+    if (auto it = cache.kernels.find(kernel_name); it != cache.kernels.end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock lock(cache.mtx);
+  if (auto it = cache.kernels.find(kernel_name); it != cache.kernels.end()) {
     return it->second;
   }
-  if (source_code.empty()) {
-    return nullptr;
-  }
-
+  std::string source_code = source_builder();
   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.
+  // Deal with long kernel names. Maximum length for filename on macOS is 255
+  // characters, and on Windows the maximum length for whole path is 260. Clip
+  // file name with a little extra room and append a 16 character hash.
+#ifdef _WIN32
+  constexpr int max_file_name_length = 140;
+#else
   constexpr int max_file_name_length = 245;
+#endif
   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);
+    auto file_id =
+        std::hash<std::string>{}(kernel_name.substr(max_file_name_length - 16));
     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";
-  auto shared_lib_path = get_temp_file(shared_lib_name.str());
+  auto output_dir = std::filesystem::temp_directory_path();
+
+  std::string shared_lib_name = "lib" + kernel_file_name + ".so";
+  auto shared_lib_path = (output_dir / shared_lib_name).string();
   bool lib_exists = false;
   {
     std::ifstream f(shared_lib_path.c_str());
@@ -81,19 +98,16 @@ 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";
-    auto source_file_path = get_temp_file(source_file_name.str());
+    std::string source_file_name = kernel_file_name + ".cpp";
+    auto source_file_path = (output_dir / source_file_name).string();
 
     std::ofstream source_file(source_file_path);
     source_file << source_code;
     source_file.close();
 
-    std::ostringstream build_command;
-    build_command << "g++ -std=c++17 -O2 -Wall -fPIC -shared "
-                  << source_file_path << " -o " << shared_lib_path;
-    std::string build_command_str = build_command.str();
-    auto return_code = system(build_command_str.c_str());
+    std::string command = JitCompiler::build_command(
+        output_dir, source_file_name, shared_lib_name);
+    auto return_code = system(command.c_str());
     if (return_code) {
       std::ostringstream msg;
       msg << "[Compile::eval_cpu] Failed to compile function " << kernel_name
@@ -103,10 +117,10 @@ void* compile(
   }
 
   // load library
-  libs.emplace_back(shared_lib_path);
+  cache.libs.emplace_back(shared_lib_path);
 
   // Load function
-  void* fun = dlsym(libs.back().lib, kernel_name.c_str());
+  void* fun = dlsym(cache.libs.back().lib, kernel_name.c_str());
   if (!fun) {
     std::ostringstream msg;
     msg << "[Compile::eval_cpu] Failed to load compiled function "
@@ -114,7 +128,7 @@ void* compile(
         << dlerror();
     throw std::runtime_error(msg.str());
   }
-  kernels.insert({kernel_name, fun});
+  cache.kernels.insert({kernel_name, fun});
   return fun;
 }
 
@@ -138,6 +152,11 @@ inline void build_kernel(
 
   NodeNamer namer;
 
+#ifdef _MSC_VER
+  // Export the symbol
+  os << "__declspec(dllexport) ";
+#endif
+
   // Start the kernel
   os << "void " << kernel_name << "(void** args) {" << std::endl;
 
@@ -266,7 +285,7 @@ void Compiled::eval_cpu(
 
   // Figure out which kernel we are using
   auto& shape = outputs[0].shape();
-  bool contiguous = compiled_check_contiguity(inputs, shape);
+  auto contiguous = compiled_check_contiguity(inputs, shape);
 
   // Handle all broadcasting and collect function input arguments
   std::vector<void*> args;
@@ -316,10 +335,7 @@ void Compiled::eval_cpu(
   }
 
   // Get the function
-  auto fn_ptr = compile(kernel_name);
-
-  // If it doesn't exist, compile it
-  if (fn_ptr == nullptr) {
+  auto fn_ptr = compile(kernel_name, [&]() {
     std::ostringstream kernel;
     kernel << get_kernel_preamble() << std::endl;
     kernel << "extern \"C\"  {" << std::endl;
@@ -334,10 +350,8 @@ void Compiled::eval_cpu(
         ndim);
     // Close extern "C"
     kernel << "}" << std::endl;
-
-    // Compile and get function pointer
-    fn_ptr = compile(kernel_name, kernel.str());
-  }
+    return kernel.str();
+  });
 
   compiled_allocate_outputs(
       inputs, outputs, inputs_, constant_ids_, contiguous, false);

mlx/backend/common/conv.cpp

@@ -3,13 +3,8 @@
 #include <cassert>
 #include <numeric>
 
-#ifdef ACCELERATE_NEW_LAPACK
-#include <Accelerate/Accelerate.h>
-#else
-#include <cblas.h>
-#endif
-
 #include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/lapack.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
@@ -684,6 +679,32 @@ void dispatch_slow_conv_3D(
 // Explicit gemm conv
 ///////////////////////////////////////////////////////////////////////////////
 
+template <typename T>
+void flip_spatial_dims_inplace(array& wt) {
+  T* x = wt.data<T>();
+  size_t out_channels = wt.shape(0);
+  size_t in_channels = wt.shape(-1);
+
+  // Calculate the total size of the spatial dimensions
+  int spatial_size = 1;
+  for (int d = 1; d < wt.ndim() - 1; ++d) {
+    spatial_size *= wt.shape(d);
+  }
+
+  for (size_t i = 0; i < out_channels; i++) {
+    T* top = x + i * spatial_size * in_channels;
+    T* bottom =
+        x + i * spatial_size * in_channels + (spatial_size - 1) * in_channels;
+    for (size_t j = 0; j < spatial_size / 2; j++) {
+      for (size_t k = 0; k < in_channels; k++) {
+        std::swap(top[k], bottom[k]);
+      }
+      top += in_channels;
+      bottom -= in_channels;
+    }
+  }
+}
+
 void explicit_gemm_conv_1D_cpu(
     const array& in,
     const array& wt,
@@ -705,7 +726,7 @@ void explicit_gemm_conv_1D_cpu(
   auto conv_dtype = float32;
 
   // Pad input
-  std::vector<int> padded_shape = {N, iH + 2 * padding[0], C};
+  Shape padded_shape = {N, iH + 2 * padding[0], C};
   array in_padded(padded_shape, conv_dtype, nullptr, {});
 
   // Fill with zeros
@@ -725,9 +746,9 @@ void explicit_gemm_conv_1D_cpu(
   copy_inplace(in, in_padded_slice, CopyType::GeneralGeneral);
 
   // Make strided view
-  std::vector<int> strided_shape = {N, oH, wH, C};
+  Shape strided_shape = {N, oH, wH, C};
 
-  std::vector<size_t> strided_strides = {
+  Strides strided_strides = {
       in_padded.strides()[0],
       in_padded.strides()[1] * wt_strides[0],
       in_padded.strides()[1],
@@ -744,7 +765,7 @@ void explicit_gemm_conv_1D_cpu(
       in_padded, strided_strides, flags, in_strided_view.size(), 0);
 
   // Materialize strided view
-  std::vector<int> strided_reshape = {N * oH, wH * C};
+  Shape strided_reshape = {N * oH, wH * C};
   array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
   copy(in_strided_view, in_strided, CopyType::General);
 
@@ -822,8 +843,7 @@ void explicit_gemm_conv_2D_cpu(
   auto conv_dtype = out.dtype();
 
   // Pad input
-  std::vector<int> padded_shape = {
-      N, iH + 2 * padding[0], iW + 2 * padding[1], C};
+  Shape padded_shape = {N, iH + 2 * padding[0], iW + 2 * padding[1], C};
   array in_padded(padded_shape, conv_dtype, nullptr, {});
 
   // Fill with zeros
@@ -844,9 +864,9 @@ void explicit_gemm_conv_2D_cpu(
   copy_inplace(in, in_padded_slice, CopyType::GeneralGeneral);
 
   // Make strided view
-  std::vector<int> strided_shape = {N, oH, oW, wH, wW, C};
+  Shape strided_shape = {N, oH, oW, wH, wW, C};
 
-  std::vector<size_t> strided_strides = {
+  Strides strided_strides = {
       in_padded.strides()[0],
       in_padded.strides()[1] * wt_strides[0],
       in_padded.strides()[2] * wt_strides[1],
@@ -860,7 +880,7 @@ void explicit_gemm_conv_2D_cpu(
       in_padded, strided_strides, flags, in_strided_view.size(), 0);
 
   // Materialize strided view
-  std::vector<int> strided_reshape = {N * oH * oW, wH * wW * C};
+  Shape strided_reshape = {N * oH * oW, wH * wW * C};
   array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
   copy(in_strided_view, in_strided, CopyType::General);
 
@@ -910,21 +930,22 @@ void explicit_gemm_conv_ND_cpu(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const bool flip) {
   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>(
+  const auto iDim =
+      Shape(in.shape().begin() + 1, in.shape().end() - 1); // Input spatial dim
+  const auto oDim = Shape(
       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
+  const auto wDim =
+      Shape(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());
+  Shape 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];
@@ -952,7 +973,7 @@ void explicit_gemm_conv_ND_cpu(
   copy_inplace(in, in_padded_slice, CopyType::GeneralGeneral);
 
   // Make strided view
-  std::vector<int> strided_shape(oDim.size() + wDim.size() + 2);
+  Shape 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];
@@ -962,7 +983,7 @@ void explicit_gemm_conv_ND_cpu(
   }
   strided_shape.back() = C;
 
-  std::vector<size_t> strided_strides(in.shape().size() * 2 - 2);
+  Strides 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];
@@ -978,7 +999,7 @@ void explicit_gemm_conv_ND_cpu(
       in_padded, strided_strides, flags, in_strided_view.size(), 0);
 
   // Materialize strided view
-  std::vector<int> strided_reshape = {N, C};
+  Shape strided_reshape = {N, C};
   for (const auto& o : oDim) {
     strided_reshape[0] *= o;
   }
@@ -1000,6 +1021,14 @@ void explicit_gemm_conv_ND_cpu(
     copy(wt, gemm_wt, ctype);
   }
 
+  if (flip) {
+    auto gemm_wt_ = array(gemm_wt.shape(), float32, nullptr, {});
+    copy(gemm_wt, gemm_wt_, CopyType::Vector);
+
+    flip_spatial_dims_inplace<float>(gemm_wt_);
+    gemm_wt = gemm_wt_;
+  }
+
   if (out.dtype() != float32) {
     gemm_out = array(out.shape(), float32, nullptr, {});
     gemm_out.set_data(allocator::malloc_or_wait(gemm_out.nbytes()));
@@ -1042,10 +1071,15 @@ void conv_1D_cpu(
     const std::vector<int>& wt_dilation,
     const std::vector<int>& in_dilation,
     bool flip) {
+  const int groups = in.shape().back() / wt.shape().back();
   if (wt_dilation[0] == 1 && in_dilation[0] == 1 && !flip) {
     return explicit_gemm_conv_1D_cpu(
         in, wt, out, padding, wt_strides, wt_dilation);
   }
+  if (wt_dilation[0] == 1 && in_dilation[0] == 1 && groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in, wt, out, padding, wt_strides, wt_dilation, flip);
+  }
 
   return dispatch_slow_conv_1D(
       in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
@@ -1060,6 +1094,13 @@ void conv_2D_cpu(
     const std::vector<int>& wt_dilation,
     const std::vector<int>& in_dilation,
     bool flip) {
+  const int groups = in.shape().back() / wt.shape().back();
+  if (wt_dilation[0] == 1 && wt_dilation[1] == 1 && in_dilation[0] == 1 &&
+      in_dilation[1] == 1 && groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in, wt, out, padding, wt_strides, wt_dilation, flip);
+  }
+
   return dispatch_slow_conv_2D(
       in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
 }
@@ -1073,6 +1114,14 @@ void conv_3D_cpu(
     const std::vector<int>& wt_dilation,
     const std::vector<int>& in_dilation,
     bool flip) {
+  const int groups = in.shape().back() / wt.shape().back();
+  if (wt_dilation[0] == 1 && wt_dilation[1] == 1 && wt_dilation[2] == 1 &&
+      in_dilation[0] == 1 && in_dilation[1] == 1 && in_dilation[2] == 1 &&
+      groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in, wt, out, padding, wt_strides, wt_dilation, flip);
+  }
+
   return dispatch_slow_conv_3D(
       in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
 }

mlx/backend/common/copy.cpp

@@ -26,292 +26,117 @@ 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) {
-  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) {
-    dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-    src_idx += i_strides[0];
+template <typename SrcT, typename DstT, int D>
+inline void copy_dims(
+    const SrcT* src,
+    DstT* dst,
+    const Shape& shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int axis) {
+  auto stride_src = i_strides[axis];
+  auto stride_dst = o_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      copy_dims<SrcT, DstT, D - 1>(
+          src, dst, shape, i_strides, o_strides, axis + 1);
+    } else {
+      *dst = static_cast<DstT>(*src);
+    }
+    src += stride_src;
+    dst += stride_dst;
   }
 }
 
 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) {
-  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) {
-      dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-      src_idx += i_strides[1];
-    }
-    src_idx += i_strides[0] - i_strides[1] * data_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) {
-  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) {
-        dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-        src_idx += i_strides[2];
-      }
-      src_idx += i_strides[1] - i_strides[2] * data_shape[2];
-    }
-    src_idx += i_strides[0] - i_strides[1] * data_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) {
-  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) {
-          dst_ptr[dst_idx++] = static_cast<DstT>(src_ptr[src_idx]);
-          src_idx += i_strides[3];
-        }
-        src_idx += i_strides[2] - i_strides[3] * data_shape[3];
-      }
-      src_idx += i_strides[1] - i_strides[2] * data_shape[2];
-    }
-    src_idx += i_strides[0] - i_strides[1] * data_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()) {
-    case 1:
-      copy_general_dim1<SrcT, DstT, stride_t>(
-          src, dst, new_shape, new_strides[0], i_offset);
-      return;
-    case 2:
-      copy_general_dim2<SrcT, DstT, stride_t>(
-          src, dst, new_shape, new_strides[0], i_offset);
-      return;
-    case 3:
-      copy_general_dim3<SrcT, DstT, stride_t>(
-          src, dst, new_shape, new_strides[0], i_offset);
-      return;
-    case 4:
-      copy_general_dim4<SrcT, DstT, stride_t>(
-          src, dst, new_shape, new_strides[0], i_offset);
-      return;
-  }
-
-  auto src_ptr = src.data<SrcT>() + i_offset;
-  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]);
-    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>
-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) {
-  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];
-    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;
-    }
-  } 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;
-    for (int i = 0; i < N; i++) {
-      *dst_ptr = static_cast<DstT>(*src_ptr);
-      src_ptr += stride_src;
-      dst_ptr += stride_dst;
-    }
-  }
-}
-
-template <typename SrcT, typename DstT, typename stride_t>
 void copy_general_general(
     const array& src,
     array& dst,
-    const std::vector<int>& data_shape,
-    const std::vector<stride_t>& i_strides,
-    const std::vector<stride_t>& o_strides,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& 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()) {
-    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);
-      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);
-      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);
-      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);
-      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);
-      return;
+  if (data_shape.empty()) {
+    auto val = static_cast<DstT>(*(src.data<SrcT>() + i_offset));
+    auto dst_ptr = dst.data<DstT>() + o_offset;
+    *dst_ptr = val;
+    return;
   }
-
-  int size = std::accumulate(
-      new_shape.end() - 5, new_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);
+  auto [shape, strides] =
+      collapse_contiguous_dims(data_shape, {i_strides, o_strides});
+  auto src_ptr = src.data<SrcT>() + i_offset;
+  auto dst_ptr = dst.data<DstT>() + o_offset;
+  int ndim = shape.size();
+  if (ndim == 1) {
+    copy_dims<SrcT, DstT, 1>(
+        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    return;
+  } else if (ndim == 2) {
+    copy_dims<SrcT, DstT, 2>(
+        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    return;
+  } else if (ndim == 3) {
+    copy_dims<SrcT, DstT, 3>(
+        src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    return;
+  }
+  ContiguousIterator in(shape, strides[0], ndim - 3);
+  ContiguousIterator out(shape, strides[1], ndim - 3);
+  auto stride = std::accumulate(
+      shape.end() - 3, shape.end(), 1, std::multiplies<int64_t>());
+  for (int64_t elem = 0; elem < src.size(); elem += stride) {
+    copy_dims<SrcT, DstT, 3>(
+        src_ptr + in.loc,
+        dst_ptr + out.loc,
+        shape,
+        strides[0],
+        strides[1],
+        ndim - 3);
+    in.step();
+    out.step();
   }
 }
 
 template <typename SrcT, typename DstT>
 inline void copy_general_general(const array& src, array& dst) {
-  return copy_general_general<SrcT, DstT, size_t>(
+  copy_general_general<SrcT, DstT>(
       src, dst, src.shape(), src.strides(), dst.strides(), 0, 0);
 }
 
+template <typename SrcT, typename DstT>
+void copy_general(
+    const array& src,
+    array& dst,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides&,
+    int64_t i_offset,
+    int64_t o_offset) {
+  copy_general_general<SrcT, DstT>(
+      src,
+      dst,
+      data_shape,
+      i_strides,
+      make_contiguous_strides(data_shape),
+      i_offset,
+      o_offset);
+}
+
+template <typename SrcT, typename DstT>
+inline void copy_general(const array& src, array& dst) {
+  copy_general_general<SrcT, DstT>(
+      src,
+      dst,
+      src.shape(),
+      src.strides(),
+      make_contiguous_strides(src.shape()),
+      0,
+      0);
+}
+
 template <typename SrcT, typename DstT, typename... Args>
 void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
   switch (ctype) {
@@ -326,6 +151,7 @@ void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
       return;
     case CopyType::GeneralGeneral:
       copy_general_general<SrcT, DstT>(src, dst, std::forward<Args>(args)...);
+      return;
   }
 }
 
@@ -426,7 +252,7 @@ inline void copy_inplace_dispatch(
 } // namespace
 
 void copy_inplace(const array& src, array& dst, CopyType ctype) {
-  return copy_inplace_dispatch(src, dst, ctype);
+  copy_inplace_dispatch(src, dst, ctype);
 }
 
 void copy(const array& src, array& dst, CopyType ctype) {
@@ -456,20 +282,19 @@ 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,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
     int64_t i_offset,
     int64_t o_offset,
     CopyType ctype) {
   switch (ctype) {
     case CopyType::General:
     case CopyType::GeneralGeneral:
-      return copy_inplace_dispatch(
+      copy_inplace_dispatch(
           src,
           dst,
           ctype,
@@ -478,31 +303,11 @@ void copy_inplace(
           o_strides,
           i_offset,
           o_offset);
-
+      break;
     case CopyType::Scalar:
     case CopyType::Vector:
-      return copy_inplace_dispatch(src, dst, ctype);
+      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

@@ -26,13 +26,12 @@ 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,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
     int64_t i_offset,
     int64_t o_offset,
     CopyType ctype);

mlx/backend/common/default_primitives.cpp

@@ -1,14 +1,10 @@
 // Copyright © 2023-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/lapack.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/primitives.h"
 
@@ -61,6 +57,7 @@ DEFAULT(Equal)
 DEFAULT(Erf)
 DEFAULT(ErfInv)
 DEFAULT(Exp)
+DEFAULT(ExpandDims)
 DEFAULT(Expm1)
 DEFAULT(FFT)
 DEFAULT(Floor)
@@ -90,7 +87,6 @@ DEFAULT_MULTI(QRF)
 DEFAULT(QuantizedMatmul)
 DEFAULT(RandomBits)
 DEFAULT(Reduce)
-DEFAULT(Reshape)
 DEFAULT(Round)
 DEFAULT(Scan)
 DEFAULT(Scatter)
@@ -105,6 +101,7 @@ DEFAULT(Softmax)
 DEFAULT(Sort)
 DEFAULT_MULTI(Split)
 DEFAULT(Square)
+DEFAULT(Squeeze)
 DEFAULT(Sqrt)
 DEFAULT(StopGradient)
 DEFAULT(Subtract)
@@ -114,6 +111,7 @@ DEFAULT(Tanh)
 DEFAULT(Transpose)
 DEFAULT(Inverse)
 DEFAULT(Cholesky)
+DEFAULT_MULTI(Eigh)
 
 namespace {
 
@@ -133,7 +131,7 @@ inline void matmul_common_general(
     } else {
       array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
       copy(arr, arr_copy, CopyType::General);
-      size_t stx = arr.shape(-1);
+      stx = arr.shape(-1);
       return std::make_tuple(false, stx, arr_copy);
     }
   };

mlx/backend/common/eigh.cpp

@@ -0,0 +1,117 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/lapack.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+void ssyevd(
+    char jobz,
+    char uplo,
+    float* a,
+    int N,
+    float* w,
+    float* work,
+    int lwork,
+    int* iwork,
+    int liwork) {
+  int info;
+  MLX_LAPACK_FUNC(ssyevd)
+  (
+      /* jobz = */ &jobz,
+      /* uplo = */ &uplo,
+      /* n = */ &N,
+      /* a = */ a,
+      /* lda = */ &N,
+      /* w = */ w,
+      /* work = */ work,
+      /* lwork = */ &lwork,
+      /* iwork = */ iwork,
+      /* liwork = */ &liwork,
+      /* info = */ &info);
+  if (info != 0) {
+    std::stringstream msg;
+    msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
+        << info;
+    throw std::runtime_error(msg.str());
+  }
+}
+
+} // namespace
+
+void Eigh::eval(const std::vector<array>& inputs, std::vector<array>& outputs) {
+  const auto& a = inputs[0];
+  auto& values = outputs[0];
+
+  auto vectors = compute_eigenvectors_
+      ? outputs[1]
+      : array(a.shape(), a.dtype(), nullptr, {});
+
+  values.set_data(allocator::malloc_or_wait(values.nbytes()));
+
+  copy(
+      a,
+      vectors,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+
+  if (compute_eigenvectors_) {
+    // Set the strides and flags so the eigenvectors
+    // are in the columns of the output
+    auto flags = vectors.flags();
+    auto strides = vectors.strides();
+    auto ndim = a.ndim();
+    std::swap(strides[ndim - 1], strides[ndim - 2]);
+
+    if (a.size() > 1) {
+      flags.row_contiguous = false;
+      if (ndim > 2) {
+        flags.col_contiguous = false;
+      } else {
+        flags.col_contiguous = true;
+      }
+    }
+    vectors.move_shared_buffer(vectors, strides, flags, vectors.data_size());
+  }
+
+  auto vec_ptr = vectors.data<float>();
+  auto eig_ptr = values.data<float>();
+
+  char jobz = compute_eigenvectors_ ? 'V' : 'N';
+  auto N = a.shape(-1);
+
+  // Work query
+  int lwork;
+  int liwork;
+  {
+    float work;
+    int iwork;
+    ssyevd(jobz, uplo_[0], nullptr, N, nullptr, &work, -1, &iwork, -1);
+    lwork = static_cast<int>(work);
+    liwork = iwork;
+  }
+
+  auto work_buf = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+  auto iwork_buf = array::Data{allocator::malloc_or_wait(sizeof(int) * liwork)};
+  for (size_t i = 0; i < a.size() / (N * N); ++i) {
+    ssyevd(
+        jobz,
+        uplo_[0],
+        vec_ptr,
+        N,
+        eig_ptr,
+        static_cast<float*>(work_buf.buffer.raw_ptr()),
+        lwork,
+        static_cast<int*>(iwork_buf.buffer.raw_ptr()),
+        liwork);
+    vec_ptr += N * N;
+    eig_ptr += N;
+  }
+}
+
+} // namespace mlx::core

mlx/backend/common/indexing.cpp

@@ -1,5 +1,4 @@
 // Copyright © 2023 Apple Inc.
-
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -33,7 +32,7 @@ void gather(
     const std::vector<array>& inds,
     array& out,
     const std::vector<int>& axes,
-    const std::vector<int>& slice_sizes) {
+    const Shape& slice_sizes) {
   // If the array is row contiguous then we can do a contiguous copy given
   // two conditions on the slice size:
   // - Any number of leading ones in the slice sizes are allowed
@@ -81,11 +80,17 @@ void gather(
   T* dst_ptr = out.data<T>();
   size_t out_idx = 0;
 
+  std::vector<ContiguousIterator> its(inds.begin(), inds.end());
+  ContiguousIterator src_it;
+  if (!can_copy && src.ndim() > 0) {
+    src_it = ContiguousIterator(slice_sizes, src.strides(), src.ndim());
+  }
   for (int idx = 0; idx < ind_size; idx++) {
     size_t src_idx = 0;
     for (int ii = 0; ii < inds.size(); ++ii) {
       auto ax = axes[ii];
-      auto idx_loc = elem_to_loc(idx, inds[ii]);
+      auto idx_loc = its[ii].loc;
+      its[ii].step();
       auto idx_val =
           offset_neg_idx(inds[ii].data<IdxT>()[idx_loc], src.shape(ax));
       src_idx += (idx_val * src.strides()[ax]);
@@ -99,9 +104,10 @@ void gather(
       out_idx += slice_size;
     } else {
       for (int jj = 0; jj < slice_size; jj++) {
-        auto src_offset = elem_to_loc(jj, slice_sizes, src.strides());
-        dst_ptr[out_idx++] = src_ptr[src_idx + src_offset];
+        dst_ptr[out_idx++] = src_ptr[src_idx + src_it.loc];
+        src_it.step();
       }
+      src_it.reset();
     }
   }
 }
@@ -112,7 +118,7 @@ void dispatch_gather(
     const std::vector<array>& inds,
     array& out,
     const std::vector<int>& axes,
-    const std::vector<int>& size) {
+    const Shape& size) {
   switch (out.dtype()) {
     case bool_:
       gather<bool, IdxT>(src, inds, out, axes, size);
@@ -216,28 +222,36 @@ void scatter(
   auto inds_ndim = updates.ndim() - out.ndim();
   size_t n_updates = nind ? inds[0].size() : 1;
 
-  std::vector<int> update_shape(
+  Shape update_shape(
       updates.shape().begin() + inds_ndim, updates.shape().end());
   size_t update_size = 1;
   for (auto us : update_shape) {
     update_size *= us;
   }
 
+  std::vector<ContiguousIterator> its(inds.begin(), inds.end());
+  ContiguousIterator update_it(updates);
+  ContiguousIterator out_it(update_shape, out.strides(), out.ndim());
+
   for (int i = 0; i < n_updates; ++i) {
     size_t out_offset = 0;
     for (int j = 0; j < nind; ++j) {
       auto ax = axes[j];
-      auto idx_loc = elem_to_loc(i, inds[j]);
+      auto idx_loc = its[j].loc;
+      its[j].step();
       auto idx_val =
           offset_neg_idx(inds[j].data<IdxT>()[idx_loc], out.shape(ax));
       out_offset += (idx_val * out.strides()[ax]);
     }
+    update_it.seek(i * update_size);
     for (int j = 0; j < update_size; ++j) {
-      auto update_loc = elem_to_loc(i * update_size + j, updates);
-      auto out_loc = elem_to_loc(j, update_shape, out.strides());
-      op(updates.data<InT>()[update_loc],
-         out.data<InT>() + out_offset + out_loc);
+      op(updates.data<InT>()[update_it.loc],
+         out.data<InT>() + out_offset + out_it.loc);
+      update_it.step();
+      out_it.step();
     }
+    out_it.reset();
+    update_it.reset();
   }
 }
 

mlx/backend/common/inverse.cpp

@@ -2,39 +2,19 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/lapack.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_(
+  MLX_LAPACK_FUNC(strtri)
+  (
       /* uplo = */ &uplo,
       /* diag = */ &diag,
       /* N = */ &N,
       /* a = */ matrix,
       /* lda = */ &N,
       /* info = */ &info);
-#endif
-
   return info;
 }
 

mlx/backend/common/jit_compiler.cpp

@@ -0,0 +1,128 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/jit_compiler.h"
+
+#include <sstream>
+#include <vector>
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+#ifdef _MSC_VER
+
+namespace {
+
+// Split string into array.
+std::vector<std::string> str_split(const std::string& str, char delimiter) {
+  std::vector<std::string> tokens;
+  std::string token;
+  std::istringstream tokenStream(str);
+  while (std::getline(tokenStream, token, delimiter)) {
+    tokens.push_back(token);
+  }
+  return tokens;
+}
+
+// Run a command and get its output.
+std::string exec(const std::string& cmd) {
+  std::unique_ptr<FILE, decltype(&_pclose)> pipe(
+      _popen(cmd.c_str(), "r"), _pclose);
+  if (!pipe) {
+    throw std::runtime_error("popen() failed.");
+  }
+  char buffer[128];
+  std::string ret;
+  while (fgets(buffer, sizeof(buffer), pipe.get())) {
+    ret += buffer;
+  }
+  // Trim trailing spaces.
+  ret.erase(
+      std::find_if(
+          ret.rbegin(),
+          ret.rend(),
+          [](unsigned char ch) { return !std::isspace(ch); })
+          .base(),
+      ret.end());
+  return ret;
+}
+
+// Get path information about MSVC.
+struct VisualStudioInfo {
+  VisualStudioInfo() {
+#ifdef _M_ARM64
+    arch = "arm64";
+#else
+    arch = "x64";
+#endif
+    // Get path of Visual Studio.
+    std::string vs_path = exec(fmt::format(
+        "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
+        " -property installationPath",
+        std::getenv("ProgramFiles(x86)")));
+    if (vs_path.empty()) {
+      throw std::runtime_error("Can not find Visual Studio.");
+    }
+    // Read the envs from vcvarsall.
+    std::string envs = exec(fmt::format(
+        "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
+        vs_path,
+        arch));
+    for (const std::string& line : str_split(envs, '\n')) {
+      // Each line is in the format "ENV_NAME=values".
+      auto pos = line.find_first_of('=');
+      if (pos == std::string::npos || pos == 0 || pos == line.size() - 1)
+        continue;
+      std::string name = line.substr(0, pos);
+      std::string value = line.substr(pos + 1);
+      if (name == "LIB") {
+        libpaths = str_split(value, ';');
+      } else if (name == "VCToolsInstallDir") {
+        cl_exe = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch);
+      }
+    }
+  }
+  std::string arch;
+  std::string cl_exe;
+  std::vector<std::string> libpaths;
+};
+
+const VisualStudioInfo& GetVisualStudioInfo() {
+  static VisualStudioInfo info;
+  return info;
+}
+
+} // namespace
+
+#endif // _MSC_VER
+
+std::string JitCompiler::build_command(
+    const std::filesystem::path& dir,
+    const std::string& source_file_name,
+    const std::string& shared_lib_name) {
+#ifdef _MSC_VER
+  const VisualStudioInfo& info = GetVisualStudioInfo();
+  std::string libpaths;
+  for (const std::string& lib : info.libpaths) {
+    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
+  }
+  return fmt::format(
+      "\""
+      "cd /D \"{0}\" && "
+      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
+      "/link /out:\"{3}\" {4} >nul"
+      "\"",
+      dir.string(),
+      info.cl_exe,
+      source_file_name,
+      shared_lib_name,
+      libpaths);
+#else
+  return fmt::format(
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared '{0}' -o '{1}'",
+      (dir / source_file_name).string(),
+      (dir / shared_lib_name).string());
+#endif
+}
+
+} // namespace mlx::core

mlx/backend/common/jit_compiler.h

@@ -0,0 +1,17 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <filesystem>
+
+namespace mlx::core {
+
+class JitCompiler {
+ public:
+  // Build a shell command that compiles a source code file to a shared library.
+  static std::string build_command(
+      const std::filesystem::path& dir,
+      const std::string& source_file_name,
+      const std::string& shared_lib_name);
+};
+
+} // namespace mlx::core

mlx/backend/common/lapack.h

@@ -1,10 +1,20 @@
-// Copyright © 2024 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
+// Required for Visual Studio.
+// https://github.com/OpenMathLib/OpenBLAS/blob/develop/docs/install.md
+#ifdef _MSC_VER
+#include <complex>
+#define LAPACK_COMPLEX_CUSTOM
+#define lapack_complex_float std::complex<float>
+#define lapack_complex_double std::complex<double>
+#endif
+
 #ifdef ACCELERATE_NEW_LAPACK
 #include <Accelerate/Accelerate.h>
 #else
+#include <cblas.h>
 #include <lapack.h>
 #endif
 

mlx/backend/common/make_compiled_preamble.ps1

@@ -0,0 +1,38 @@
+# This script generates a C++ function that provides the CPU
+# code for use with kernel generation.
+#
+# Copyright © 2024 Apple Inc.
+
+$OUTPUT_FILE = $args[0]
+$CL = $args[1]
+$SRCDIR = $args[2]
+
+# Get command result as array.
+$CONTENT = & $CL /std:c++17 /EP "/I$SRCDIR" /Tp "$SRCDIR/mlx/backend/common/compiled_preamble.h"
+# Remove empty lines.
+# Otherwise there will be too much empty lines making the result unreadable.
+$CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
+# Concatenate to string.
+$CONTENT = $CONTENT -join "`n"
+
+# Append extra content.
+$CONTENT = @"
+$($CONTENT)
+using namespace mlx::core;
+using namespace mlx::core::detail;
+"@
+
+# Convert each char to ASCII code.
+# Unlike the unix script that outputs string literal directly, the output from
+# MSVC is way too large to be embedded as string and compilation will fail, so
+# we store it as static array instead.
+$CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
+
+$OUTPUT = @"
+const char* get_kernel_preamble() {
+  static char preamble[] = { $CHARCODES };
+  return preamble;
+}
+"@
+
+Set-Content -Path $OUTPUT_FILE -Value $OUTPUT

mlx/backend/common/make_compiled_preamble.sh

@@ -10,18 +10,21 @@ OUTPUT_FILE=$1
 GCC=$2
 SRCDIR=$3
 CLANG=$4
+ARCH=$5
 
 if [ "$CLANG" = "TRUE" ]; then
   read -r -d '' INCLUDES <<- EOM
-  #include <cmath>
-  #include <complex>
-  #include <cstdint>
-  #include <vector>
+#include <cmath>
+#include <complex>
+#include <cstdint>
+#include <vector>
 EOM
-
+CC_FLAGS="-arch ${ARCH}"
+else
+CC_FLAGS="-std=c++17"
 fi
 
-CONTENT=$($GCC -I "$SRCDIR" -E "$SRCDIR/mlx/backend/common/compiled_preamble.h" 2>/dev/null)
+CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E "$SRCDIR/mlx/backend/common/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
 const char* get_kernel_preamble() {

mlx/backend/common/masked_mm.cpp

@@ -1,15 +1,10 @@
 // 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/lapack.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/primitives.h"
 
@@ -24,10 +19,10 @@ inline void mask_matrix(
     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 int64_t X_data_str,
+    const int64_t Y_data_str,
+    const int64_t X_mask_str,
+    const int64_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;
@@ -89,7 +84,7 @@ void BlockMaskedMM::eval(const std::vector<array>& inputs, array& out) {
         } else {
           array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
           copy(arr, arr_copy, CopyType::General);
-          size_t stx = arr.shape(-1);
+          int64_t stx = arr.shape(-1);
           return std::make_tuple(false, stx, arr_copy);
         }
       };
@@ -122,13 +117,13 @@ void BlockMaskedMM::eval(const std::vector<array>& inputs, array& out) {
                        int Y,
                        size_t X_data_str,
                        size_t Y_data_str) {
-    size_t mask_offset = elem_to_loc(
+    auto 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];
+    auto X_mask_str = mask.strides()[mask.ndim() - 2];
+    auto Y_mask_str = mask.strides()[mask.ndim() - 1];
 
     if (mask.dtype() == bool_) {
       return mask_matrix(
@@ -235,7 +230,7 @@ void GatherMM::eval(const std::vector<array>& inputs, array& out) {
     } else {
       array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
       copy(arr, arr_copy, CopyType::General);
-      size_t stx = arr.shape(-1);
+      int64_t stx = arr.shape(-1);
       return std::make_tuple(false, stx, arr_copy);
     }
   };
@@ -267,13 +262,13 @@ void GatherMM::eval(const std::vector<array>& inputs, array& out) {
   auto& lhs_indices = inputs[2];
   auto& rhs_indices = inputs[3];
 
-  std::vector<int> batch_shape = get_batch_dims(out.shape());
+  auto 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());
+  auto batch_shape_A = get_batch_dims(a.shape());
+  auto batch_strides_A = get_batch_dims(a.strides());
+  auto batch_shape_B = get_batch_dims(b.shape());
+  auto 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>();

mlx/backend/common/ops.h

@@ -295,6 +295,13 @@ struct Floor {
   }
 };
 
+struct Imag {
+  template <typename T>
+  T operator()(T x) {
+    return std::imag(x);
+  }
+};
+
 struct Log {
   template <typename T>
   T operator()(T x) {
@@ -337,6 +344,13 @@ struct Negative {
   }
 };
 
+struct Real {
+  template <typename T>
+  T operator()(T x) {
+    return std::real(x);
+  }
+};
+
 struct Round {
   template <typename T>
   T operator()(T x) {
@@ -486,7 +500,12 @@ struct Equal {
 struct NaNEqual {
   template <typename T>
   bool operator()(T x, T y) {
-    return x == y || (std::isnan(x) && std::isnan(y));
+    if constexpr (std::is_integral_v<T>) {
+      // isnan always returns false for integers, and MSVC refuses to compile.
+      return x == y;
+    } else {
+      return x == y || (std::isnan(x) && std::isnan(y));
+    }
   }
 };