fix test

cpu fallback
add checks
2024-12-06 10:26:54 -08:00 · 2024-12-06 01:22:50 -08:00 · 2024-12-06 01:09:00 -08:00 · 2024-12-06 00:45:01 -08:00 · 2024-12-06 00:21:05 -08:00
865 changed files with 27746 additions and 112840 deletions
@@ -0,0 +1,413 @@
+version: 2.1
+
+orbs:
+  apple: ml-explore/pr-approval@0.1.0
+
+parameters:
+  nightly_build:
+    type: boolean
+    default: false
+  weekly_build:
+    type: boolean
+    default: false
+  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
+
+    steps:
+      - checkout
+      - run:
+          name: Run style checks
+          command: |
+            pip install pre-commit
+            pre-commit run --all
+            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
+      - run:
+          name: Install dependencies
+          command: |
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install numpy
+            sudo apt-get update
+            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
+      - run:
+          name: Install Python package
+          command: |
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              python3 setup.py build_ext --inplace
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              python3 setup.py develop
+      - run:
+          name: Generate package stubs
+          command: |
+            echo "stubs"
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Run Python tests
+          command: |
+            python3 -m unittest discover python/tests -v
+      - run:
+          name: Build CPP only
+          command: |
+            mkdir -p build && cd build 
+            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
+            make -j `nproc`
+      - run:
+          name: Run CPP tests
+          command: ./build/tests/tests
+
+  mac_build_and_test:
+    parameters:
+      xcode_version:
+        type: string
+        default: "15.2.0"
+    macos:
+      xcode: << parameters.xcode_version >>
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            brew install python@3.9
+            brew install openmpi
+            python3.9 -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install numpy
+            pip install torch
+            pip install tensorflow
+            pip install unittest-xml-reporting
+      - run:
+          name: Install Python package
+          command: |
+            source env/bin/activate
+            DEBUG=1 CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` pip install -e . -v
+      - run:
+          name: Generate package stubs
+          command: |
+            source env/bin/activate
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Run Python tests
+          command: |
+            source env/bin/activate
+            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
+            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
+            mpirun --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: |
+            source env/bin/activate
+            cd examples/extensions
+            pip install -r requirements.txt
+            python setup.py build_ext -j8
+      - store_test_results:
+          path: test-results
+      - run:
+          name: Build CPP only
+          command: |
+            source env/bin/activate
+            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
+      - run:
+          name: Run CPP tests
+          command: |
+            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
+      - run:
+          name: Build small binary
+          command: |
+            source env/bin/activate
+            cd build/
+            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
+              -DBUILD_SHARED_LIBS=ON \
+              -DMLX_BUILD_CPU=OFF \
+              -DMLX_BUILD_SAFETENSORS=OFF \
+              -DMLX_BUILD_GGUF=OFF \
+              -DMLX_METAL_JIT=ON
+            make -j `sysctl -n hw.ncpu`
+      - run:
+          name: Run Python tests with JIT
+          command: |
+            source env/bin/activate
+            CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
+              pip install -e . -v
+            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_jit
+
+  build_release:
+    parameters:
+      python_version:
+        type: string
+        default: "3.9"
+      xcode_version:
+        type: string
+        default: "15.2.0"
+      build_env:
+        type: string
+        default: ""
+    macos:
+      xcode: << parameters.xcode_version >>
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            brew install python@<< parameters.python_version >>
+            brew install openmpi
+            python<< parameters.python_version >> -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install --upgrade setuptools
+            pip install numpy
+            pip install twine
+            pip install build
+      - run:
+          name: Install Python package
+          command: |
+            source env/bin/activate
+            DEV_RELEASE=1 \
+              CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              pip install . -v
+      - run:
+          name: Generate package stubs
+          command: |
+            source env/bin/activate
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Build Python package
+          command: |
+            source env/bin/activate
+            << parameters.build_env >> \
+              CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              python -m build -w
+      - when:
+          condition: << parameters.build_env >>
+          steps:
+            - run:
+                name: Upload package
+                command: |
+                  source env/bin/activate
+                  twine upload dist/*
+      - store_artifacts:
+          path: dist/
+
+  build_linux_release:
+    parameters:
+      python_version:
+        type: string
+        default: "3.9"
+      extra_env:
+        type: string
+        default: "DEV_RELEASE=1"
+    docker:
+      - image: ubuntu:20.04
+    steps:
+      - checkout
+      - run:
+          name: Build wheel
+          command: |
+            PYTHON=python<< parameters.python_version >>
+            apt-get update
+            apt-get upgrade -y
+            DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
+            apt-get install -y apt-utils
+            apt-get install -y software-properties-common
+            add-apt-repository -y ppa:deadsnakes/ppa
+            apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
+            apt-get install -y libblas-dev liblapack-dev liblapacke-dev
+            apt-get install -y build-essential git
+            $PYTHON -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.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
+            pip install typing_extensions
+            python setup.py generate_stubs 
+            << parameters.extra_env >> \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              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/
+
+workflows:
+  build_and_test:
+    when:
+      and:
+        - matches:
+            pattern: "^(?!pull/)[-\\w]+$"
+            value: << pipeline.git.branch >>
+        - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.weekly_build >>
+        - not: << pipeline.parameters.test_release >>
+    jobs:
+      - mac_build_and_test:
+          matrix:
+            parameters:
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
+      - linux_build_and_test
+      - build_documentation 
+
+  build_pypi_release:
+    when:
+      and:
+        - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.weekly_build >>
+        - not: << pipeline.parameters.test_release >>
+    jobs:
+      - build_release:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+          matrix:
+            parameters:
+              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:
+        pattern: "^pull/\\d+(/head)?$"
+        value: << pipeline.git.branch >>
+    jobs:
+      - hold:
+          type: approval
+      - apple/authenticate:
+          context: pr-approval
+      - mac_build_and_test:
+          requires: [ hold ]
+          matrix:
+            parameters:
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
+      - linux_build_and_test:
+          requires: [ hold ]
+  nightly_build:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.nightly_build >>
+    jobs:
+      - build_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              xcode_version: ["15.0.0", "15.2.0"]
+  weekly_build:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.weekly_build >>
+    jobs:
+      - build_release:
+          matrix:
+            parameters:
+              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.linux_release >>
+    jobs:
+      - build_linux_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              extra_env: ["PYPI_RELEASE=1"]
@@ -1,31 +0,0 @@
-name: 'Build CUDA wheel'
-description: 'Build CUDA wheel'
-
-inputs:
-  arch:
-    description: 'Platform architecture tag'
-    required: true
-    type: choice
-    options:
-      - x86_64
-      - aarch64
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build package
-      shell: bash
-      env:
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
-      run: |
-        pip install auditwheel "build<=1.4.2" patchelf setuptools
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-
-        auditwheel repair dist/mlx_cuda*.whl \
-          --plat manylinux_2_35_${{ inputs.arch }} \
-          --exclude libcublas* \
-          --exclude libcuda* \
-          --exclude libcudnn* \
-          --exclude libnccl* \
-          --exclude libnvrtc*
@@ -1,38 +0,0 @@
-name: 'Build Documentation'
-description: 'Build documentation'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Setup machine
-      uses: ./.github/actions/setup-linux
-
-    - name: Install dependencies
-      shell: bash
-      run: |
-        sudo apt-get install -y doxygen
-        source .venv/bin/activate
-        pip install -r docs/requirements.txt
-        pip install . -v
-  
-    - name: Build documentation
-      shell: bash
-      run: |
-        source .venv/bin/activate
-        cd docs
-        doxygen
-        make html O=-W
-    
-    - name: Create artifact tar
-      shell: bash
-      run: tar -cf artifact.tar -C docs --dereference build/html index.html
-
-    # Do it manually because upload-pages-artifact requires gtar
-    - name: Upload artifact
-      id: upload-artifact
-      uses: actions/upload-artifact@v5
-      with:
-        name: github-pages
-        path: artifact.tar
-        retention-days: 1
-        if-no-files-found: error
@@ -1,42 +0,0 @@
-name: 'Build Linux wheel'
-description: 'Build Linux wheel'
-
-inputs:
-  build-backend:
-    description: 'Build the backend mlx-cpu package'
-    type: boolean
-    required: false
-    default: false
-  arch:
-    description: 'Platform architecture tag'
-    required: true
-    type: choice
-    options:
-      - x86_64
-      - aarch64
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build MLX
-      shell: bash
-      run: pip install -e . -v
-
-    - name: Build Python package
-      shell: bash
-      run: |
-        pip install auditwheel patchelf "build<=1.4.2"
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        auditwheel repair dist/mlx-*.whl \
-          --plat manylinux_2_35_${{ inputs.arch }} \
-          --exclude libmlx.so* \
-          --only-plat
-
-    - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash
-      run: |
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_${{ inputs.arch }}
@@ -1,38 +0,0 @@
-name: 'Build and Test on Linux'
-
-inputs:
-  toolkit:
-    description: 'The toolkit to build with'
-    required: false
-    default: 'cpu'
-
-runs:
-  using: "composite"
-  steps:
-
-    - name: Install Python package
-      id: python_build
-      shell: sh
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: >-
-          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
-          -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }}
-      run: |
-        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
-          # There is no GPU in arm64 runner, use a common arch.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
-          # Can not build tests and stubs when the built executables can not run.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
-        fi
-        # Install cpu-only torch to save space
-        pip install torch --index-url https://download.pytorch.org/whl/cpu
-        pip install --no-build-isolation -e ".[dev]" -v
-        # Pass the CMAKE_ARGS to following steps.
-        echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
-
-    - name: Build CPP only
-      shell: bash
-      run: |
-        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }}
-        cmake --build build -j $(nproc)
@@ -1,36 +0,0 @@
-name: 'Build macOS release'
-description: 'Build MLX releases macOS'
-
-inputs:
-  macos-target:
-    description: 'macOS build target'
-    required: false
-    default: '15.0'
-  build-backend:
-    description: 'Build the backend mlx-metal package'
-    type: boolean
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build Python package
-      shell: bash -l {0}
-      env:
-        DEVELOPER_DIR: /Applications/Xcode-latest.app
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        pip install build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-
-    - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash -l {0}
-      env:
-        DEVELOPER_DIR: /Applications/Xcode-latest.app
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
@@ -1,82 +0,0 @@
-name: 'Build and Test on macOS'
-description: 'Build and test MLX on macOS'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install dependencies
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
-      run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
-        pip install -e ".[dev]" -v
-
-    - name: Install tests dependencies
-      shell: bash -l {0}
-      run: |
-        pip install tensorflow
-
-    - name: Run Python tests
-      shell: bash -l {0}
-      env:
-        LOW_MEMORY: 1
-      run: |
-        DEVICE=cpu python -m unittest discover -v python/tests
-        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
-        mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
-        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-        if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-    
-    - name: Build example extension
-      shell: bash -l {0}
-      run: |
-        cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
-    - name: Build CPP only
-      shell: bash -l {0}
-      run: |
-        mkdir -p build
-        cd build
-        cmake ..
-        make -j $(sysctl -n hw.ncpu)
-
-    - name: Run CPP tests
-      shell: bash -l {0}
-      env:
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: |
-        ./build/tests/tests
-        ./build/tests/test_teardown
-
-    - name: Build small binary with JIT
-      shell: bash -l {0}
-      run: |
-        mkdir -p build
-        cd build
-        cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
-          -DBUILD_SHARED_LIBS=ON \
-          -DMLX_BUILD_CPU=OFF \
-          -DMLX_BUILD_SAFETENSORS=OFF \
-          -DMLX_BUILD_GGUF=OFF \
-          -DMLX_METAL_JIT=ON
-        make -j $(sysctl -n hw.ncpu)
-    
-    - name: Run Python tests with JIT
-      shell: bash -l {0}
-      env:
-        LOW_MEMORY: 1
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: |
-        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
-        python -m unittest discover -v python/tests
@@ -1,26 +0,0 @@
-name: 'Build on Windows'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Install Python package
-      id: python-build
-      shell: cmd
-      env:
-        # For MSVC, Ninja/Release is the only config supported by ccache.
-        CMAKE_ARGS: >-
-          -G Ninja
-          -DCMAKE_BUILD_TYPE=Release
-          -DCMAKE_C_COMPILER=cl
-          -DCMAKE_CXX_COMPILER=cl
-          -DCMAKE_RC_COMPILER=rc
-      run: |
-        uv pip install ".[dev]" -v
-        :: Pass the CMAKE_ARGS to following steps.
-        >>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS%
-
-    - name: Build CPP only
-      shell: cmd
-      run: |
-        cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }}
-        cmake --build build -j %NUMBER_OF_PROCESSORS%
@@ -1,102 +0,0 @@
-name: 'Setup Linux Environment'
-description: 'Install dependencies for Linux builds'
-
-inputs:
-  toolkit:
-    description: 'Which toolkit to install'
-    required: false
-    default: 'cpu'
-  python-version:
-    description: 'Version of python to set up'
-    required: false
-    default: '3.14'
-  use-ccache:
-    description: 'Whether to enable ccache'
-    required: false
-    default: 'true'
-  ccache-key:
-    required: false
-    default: 'ccache'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install common dependencies
-      shell: bash
-      run: |
-        echo "::group::Install common dependencies"
-        sudo apt-get update
-        sudo apt-get install -y --no-install-recommends \
-            zip \
-            libblas-dev liblapack-dev liblapacke-dev \
-            openmpi-bin openmpi-common libopenmpi-dev
-        echo "::endgroup::"
-
-    - name: Use ccache
-      if: ${{ inputs.use-ccache == 'true' }}
-      uses: hendrikmuhs/ccache-action@v1.2
-      with:
-        key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
-        max-size: 1GB
-        # ccache-action bug: running "apt-get update" fails on large arm runner.
-        update-package-index: false
-
-    - uses: actions/setup-python@v6
-      with:
-        python-version: ${{ inputs.python-version }}
-
-    - name: Setup Python venv
-      shell: bash
-      run: |
-        echo "::group::Setup Python venv"
-        python -m venv .venv
-        source .venv/bin/activate
-        pip install setuptools cmake typing_extensions
-        echo PATH=$PATH >> $GITHUB_ENV
-        # Search python packages in .venv
-        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
-        echo "::endgroup::"
-
-    - name: Set swap space
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
-      with:
-        swap-size-gb: 16
-
-    - name: Install CUDA toolkit
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      shell: bash
-      env:
-        # Note: the CI machine does not meet CUDA 13's driver requirement.
-        # Compatibility matrix:
-        # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
-        PACKAGES: |
-          {
-            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6",
-            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9",
-            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0"
-          }
-      run: |
-        echo "::group::Install CUDA toolkit"
-        # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is
-        # Jetson specific. SBSA means Arm Server Base System Architecture.
-        ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
-        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
-        sudo dpkg -i cuda-keyring_1.1-1_all.deb
-        sudo apt-get update
-        sudo apt-get install -y --no-install-recommends \
-            libnccl2 libnccl-dev \
-            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
-        echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH
-        echo "::endgroup::"
-
-    - name: CUDA packages and driver report
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      shell: bash
-      run: |
-        echo "::group::Installed NVIDIA and CUDA packages"
-        dpkg -l | egrep "cuda|nvidia" -i
-        echo "::endgroup::"
-        echo "::group::NVIDIA-SMI Status"
-        nvidia-smi || true
-        echo "::endgroup::"
@@ -1,24 +0,0 @@
-name: 'Setup macOS Environment'
-description: 'Install dependencies for macOS builds'
-
-inputs:
-  python-version:
-    description: 'Python version to use'
-    required: false
-    default: '3.10'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Homebrew packages
-      shell: sh
-      run: /opt/homebrew/bin/brew install openmpi
-    
-    - name: Verify MetalToolchain installed
-      shell: bash
-      run: xcodebuild -showComponent MetalToolchain
-
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
@@ -1,42 +0,0 @@
-name: 'Setup Windows environment'
-
-inputs:
-  python-version:
-    description: 'Version of python to set up'
-    required: false
-    default: '3.14'
-  use-ccache:
-    description: 'Whether to enable ccache'
-    required: false
-    default: 'true'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Use ccache
-      if: ${{ inputs.use-ccache == 'true' }}
-      uses: hendrikmuhs/ccache-action@v1.2
-      with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-cpu
-        max-size: 1GB
-
-    - name: Setup Visual Studio cmd
-      shell: cmd
-      run: |
-        :: Find out path to VS.
-        pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\"
-        for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x
-        popd
-        :: Import VS vars.
-        call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" x64
-        :: Export to all steps.
-        >>%GITHUB_ENV% set
-
-    - uses: astral-sh/setup-uv@v7
-
-    - name: Setup Python venv
-      shell: cmd
-      run: |
-        uv venv --python ${{ inputs.python-version }}
-        call ".venv/Scripts/activate.bat"
-        >>%GITHUB_ENV% set
@@ -1,69 +0,0 @@
-name: 'Run Linux tests'
-
-inputs:
-  has-gpu:
-    description: 'Run GPU tests'
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Run MPI tests
-      shell: bash
-      run: |
-        echo "::group::MPI tests"
-        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-        echo "::endgroup::"
-
-    - name: Run distributed tests
-      if: ${{ inputs.has-gpu == 'false' }}
-      shell: bash
-      run: |
-        echo "::group::Distributed tests"
-        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-        if grep -Fq '[WARN]' stderr.log ; then
-          grep -F '[WARN]' stderr.log
-          echo "Distributed ring test failed";
-          exit 1;
-        fi
-        echo "::endgroup::"
-
-    - name: Run Python tests - CPU
-      if: ${{ inputs.has-gpu == 'false' }}
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::Python tests - CPU"
-        python -m unittest discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run Python tests - GPU
-      if: ${{ inputs.has-gpu == 'true' }}
-      shell: bash
-      env:
-        DEVICE: gpu
-      run: |
-        echo "::group::Python tests - GPU"
-        python -m tests discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run CPP tests - CPU
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::CPP tests - CPU"
-        ./build/tests/tests
-        echo "::endgroup::"
-
-    - name: Run CPP tests - GPU
-      if: ${{ inputs.has-gpu == 'true' }}
-      shell: bash
-      env:
-        DEVICE: gpu
-      run: |
-        echo "::group::CPP tests - GPU"
-        ./build/tests/tests -sfe="*linalg_tests.cpp"
-        echo "::endgroup::"
@@ -1,21 +0,0 @@
-name: 'Run tests on Windows'
-
-runs:
-  using: 'composite'
-  steps:
-    - name: Run Python tests - CPU
-      shell: bash
-      run: |
-        echo "::group::Python tests - CPU"
-        python -m unittest discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run CPP tests - CPU
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::CPP tests - CPU"
-        ./build/tests.exe -tce="*gguf*,test random uniform"
-        ./build/test_teardown.exe
-        echo "::endgroup::"
@@ -1,6 +0,0 @@
-version: 2
-updates:
-  - package-ecosystem: "github-actions"
-    directory: "/"
-    schedule:
-      interval: "weekly"
@@ -1,48 +0,0 @@
-#!/bin/bash
-set -ex
-
-export CMAKE_C_COMPILER=/usr/bin/clang
-export CMAKE_CXX_COMPILER=/usr/bin/clang++
-BASE_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=DEBUG -DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-if [[ "$(uname -s)" != "Darwin" ]]; then
-  BASE_CMAKE_ARGS+=" -DMLX_BUILD_METAL=OFF"
-fi
-
-run_test() {
-  local sanitizer_name=$1
-  local cmake_sanitizer_flag="-DUSE_${sanitizer_name}=ON"
-  echo "  Running tests with: ${sanitizer_name}"
-
-  case "$sanitizer_name" in
-    ASAN)
-      export ASAN_OPTIONS="detect_leaks=0"
-      ;;
-    UBSAN)
-      export UBSAN_OPTIONS="halt_on_error=0:print_stacktrace=1"
-      ;;
-    TSAN)
-      export TSAN_OPTIONS=""
-      ;;
-  esac
-
-  rm -rf build
-  mkdir -p build
-  pushd build > /dev/null
-
-  cmake .. ${BASE_CMAKE_ARGS} ${cmake_sanitizer_flag}
-  make -j $(nproc)
-  ./tests/tests
-
-  popd > /dev/null
-  unset ${sanitizer_name}_OPTIONS
-}
-
-sanitizer_arg=$(echo "$1" | tr '[:lower:]' '[:upper:]')
-
-if [[ "$sanitizer_arg" == "ASAN" || "$sanitizer_arg" == "UBSAN" || "$sanitizer_arg" == "TSAN" ]]; then
-  run_test "$sanitizer_arg"
-  echo "  ${sanitizer_arg} test run completed successfully."
-else
-  echo "Error: Invalid sanitizer '$1'. Please use one of: ASAN, UBSAN, TSAN."
-  exit 1
-fi
@@ -1,27 +0,0 @@
-#!/bin/bash
-set -ex
-
-# [Setup] Install dependencies inside the container.
-dnf update -y
-dnf install -y \
-  blas-devel \
-  lapack-devel \
-  openblas-devel \
-  make \
-  cmake \
-  clang \
-  git
-dnf clean all
-
-# [C++] CI Build Sanity Check: Verifies code compilation, not for release.
-export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-export DEBUG=1
-export CMAKE_C_COMPILER=/usr/bin/clang
-export CMAKE_CXX_COMPILER=/usr/bin/clang++
-
-mkdir -p build
-pushd build
-cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-make -j $(nproc)
-./tests/tests
-popd
@@ -1,152 +0,0 @@
-name: Build and Test
-
-on:
-  pull_request:
-  push:
-    branches:
-      - main
-      # For testing CI without starting a pull request:
-      - test/*
-
-permissions:
-  contents: read
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
-
-jobs:
-  check_lint:
-    name: Check Lint
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: pre-commit/action@v3.0.1
-
-  linux_build_and_test:
-    name: Linux (cpu, ${{ matrix.arch }})
-    needs: check_lint
-    strategy:
-      fail-fast: false
-      matrix:
-        arch: ['x86_64', 'aarch64']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
-      - uses: ./.github/actions/test-linux
-      - run: df -h
-
-  cuda_build_and_test:
-    name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }})
-    if: github.repository == 'ml-explore/mlx'
-    needs: check_lint
-    strategy:
-      fail-fast: false
-      matrix:
-        arch: ['x86_64', 'aarch64']
-        toolkit: ['cuda-12.6', 'cuda-12.9']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'gpu-t4-4-core' || 'ubuntu-22.04-arm' }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/build-linux
-        with:
-          toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/test-linux
-        if: matrix.arch == 'x86_64'
-        with:
-          has-gpu: true
-
-  mac_build_and_test:
-    name: macOS (${{ matrix.macos-target }})
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        macos-target: ["14.0", "15.0", "26.0"]
-    runs-on: [self-hosted, macos]
-    env:
-      MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }}
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-macos
-      - uses: ./.github/actions/build-macos
-
-  windows_build_and_test:
-    name: Windows (cpu, x86_64)
-    needs: check_lint
-    runs-on: windows-2025
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-windows
-      - uses: ./.github/actions/build-windows
-      - uses: ./.github/actions/test-windows
-
-  build_documentation:
-    name: Build Documentation
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22.04
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/build-docs
-
-  linux_sanitizer_build_and_test:
-    name: Linux Sanitizer Tests (${{ matrix.sanitizer }})
-    needs: check_lint
-    strategy:
-      fail-fast: false
-      matrix:
-        sanitizer: [ASAN, UBSAN]
-        # todo 12/16/2025: enable TSAN later + consider enabling ASAN for GPU backend tests.
-        # sanitizer: [ASAN, UBSAN, TSAN]
-    runs-on: ubuntu-22.04-arm
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v6
-
-      - name: Install Dependencies
-        run: |
-          export DEBIAN_FRONTEND=noninteractive
-          sudo apt-get update -y
-          sudo apt-get install -y \
-            build-essential \
-            libblas-dev \
-            liblapacke-dev \
-            libopenblas-dev \
-            cmake \
-            clang \
-            git
-          sudo apt-get clean
-          sudo rm -rf /var/lib/apt/lists/*
-
-      - name: Linux Build and Test with ${{ matrix.sanitizer }}
-        run: |
-          bash .github/scripts/build-sanitizer-tests.sh ${{ matrix.sanitizer }}
-
-  linux_fedora_build_cpp:
-    name: Linux Fedora (${{ matrix.arch }})
-    needs: check_lint
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - host: ubuntu-22.04
-            arch: x86_64
-          - host: ubuntu-22.04-arm
-            arch: aarch64
-
-    runs-on: ${{ matrix.host }}
-    container:
-      image: fedora:42
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v6
-
-      - name: CPP Build Test - No Release
-        run: |
-          bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh
@@ -1,28 +0,0 @@
-name: Documentation
-
-on:
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build:
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/build-docs
-      
-  deploy:
-    needs: build
-    permissions:
-      pages: write
-      id-token: write
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v5
@@ -1,108 +0,0 @@
-name: Nightly Build
-
-on:
-  schedule:
-    - cron: 33 6 * * 1-5
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build_linux_release:
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.10", "3.14"]
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux-release
-        with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
-          arch: "x86_64"
-      - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v7
-        with:
-          name: linux-wheels-${{ matrix.python_version }}
-          path: wheelhouse/mlx-*.whl
-          retention-days: 7
-      - name: Upload mlx-cpu artifacts
-        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v7
-        with:
-          name: mlx-cpu
-          path: wheelhouse/mlx_cpu-*.whl
-          retention-days: 7
-      - run: df -h
-
-  build_linux_with_tests:
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.11", "3.12", "3.13", "3.14"]
-        runner:
-          - ubuntu-22.04
-          - ubuntu-22.04-arm
-    runs-on: ${{ matrix.runner }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux
-      - uses: ./.github/actions/test-linux
-      - run: df -h
-
-  build_mac_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python-version: ["3.10", "3.13"]
-    runs-on: [self-hosted, macos]
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-macos
-        with:
-          python-version: ${{ matrix.python-version }}
-      - uses: ./.github/actions/build-macos
-      - name: Build macOS 26 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 26.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 15 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 15.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 14 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 14.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-
-  build_cuda_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        arch: ['x86_64', 'aarch64']
-        toolkit: ['cuda-12.9', 'cuda-13.0']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-        with:
-          arch: ${{ matrix.arch }}
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v7
-        with:
-          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
-          path: wheelhouse/mlx_cuda_*.whl
-          retention-days: 7
@@ -0,0 +1,20 @@
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check_lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pre-commit black isort clang-format
+      - name: Run lint
+        run: |
+          pre-commit run --all-files
@@ -1,256 +0,0 @@
-name: PyPI Release
-
-on:
-  push:
-    tags:
-      - 'v*'
-    branches:
-      - 'test-publish/*'
-  workflow_dispatch:
-    inputs:
-      dry_run:
-        description: 'Dry run (do not publish to PyPi)'
-        required: false
-        type: boolean
-      dev_release:
-        description: 'Development release (DEV_RELEASE=1)'
-        required: false
-        type: boolean
-
-permissions:
-  contents: read
-
-jobs:
-  build_documentation:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/build-docs
-
-  deploy_documentation:
-    if: ${{ !inputs.dry_run }}
-    needs: build_documentation
-    permissions:
-      pages: write
-      id-token: write
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v5
-
-  build_linux_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-        arch: ['x86_64', 'aarch64']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
-    env:
-      PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          python-version: ${{ matrix.python_version }}
-          use-ccache: false
-      - uses: ./.github/actions/build-linux-release
-        with:
-          build-backend: ${{ matrix.python_version == '3.10' }}
-          arch: ${{ matrix.arch }}
-      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v7
-        with:
-          overwrite: true
-          name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
-          path: wheelhouse/mlx-*.whl
-          if-no-files-found: error
-      - name: Upload CPU artifacts
-        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v7
-        with:
-          overwrite: true
-          name: mlx-cpu-${{ matrix.arch }}
-          path: wheelhouse/mlx_cpu-*.whl
-          if-no-files-found: error
-
-  build_mac_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: [self-hosted, macos]
-    env:
-      PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-macos
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools typing_extensions
-          pip install -e . -v
-      - name: Build macOS 14 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 14.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 15 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 15.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Build macOS 26 package
-        uses: ./.github/actions/build-macos-release
-        with:
-          macos-target: 26.0
-          build-backend: ${{ matrix.python-version == '3.10' }}
-      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v7
-        with:
-          overwrite: true
-          name: mac-wheels-${{ matrix.python-version }}
-          path: dist/mlx-*.whl
-          if-no-files-found: error
-      - name: Upload Metal artifacts
-        if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v7
-        with:
-          overwrite: true
-          name: mlx-metal
-          path: dist/mlx_metal-*.whl
-          if-no-files-found: error
-
-  build_cuda_release:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      matrix:
-        arch: ['x86_64', 'aarch64']
-        toolkit: ['cuda-12.9', 'cuda-13.0']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
-    env:
-      PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: ${{ matrix.toolkit }}
-          ccache-key: 'ccache-release'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-        with:
-          arch: ${{ matrix.arch }}
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v7
-        with:
-          overwrite: true
-          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
-          path: wheelhouse/mlx_cuda_*.whl
-          if-no-files-found: error
-
-  pypi-publish:
-    name: Upload release to PyPI
-    runs-on: ubuntu-latest
-    needs: [build_linux_release, build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
-      url: https://pypi.org/p/mlx
-    steps:
-      - uses: actions/download-artifact@v8
-        with:
-          pattern: linux-wheels-*
-          merge-multiple: true
-          path: dist
-      - uses: actions/download-artifact@v8
-        with:
-          pattern: mac-wheels-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: du -ah dist
-      - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: https://upload.pypi.org/legacy/
-
-  pypi-publish-cuda:
-    name: Upload CUDA release to PyPI
-    runs-on: ubuntu-latest
-    needs: [build_cuda_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
-      url: https://pypi.org/p/mlx-cuda
-    steps:
-      - uses: actions/download-artifact@v8
-        with:
-          pattern: mlx-cuda-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: du -ah dist
-      - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: https://upload.pypi.org/legacy/
-
-  pypi-publish-cpu:
-    name: Upload CPU release to PyPI
-    runs-on: ubuntu-latest
-    needs: [build_linux_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
-      url: https://pypi.org/p/mlx-cpu
-    steps:
-      - uses: actions/download-artifact@v8
-        with:
-          pattern: mlx-cpu-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: du -ah dist
-      - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: https://upload.pypi.org/legacy/
-
-  pypi-publish-metal:
-    name: Upload Metal release to PyPI
-    runs-on: ubuntu-latest
-    needs: [build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
-      url: https://pypi.org/p/mlx-metal
-    steps:
-      - uses: actions/download-artifact@v8
-        with:
-          name: mlx-metal
-          path: dist
-      - name: Display structure of downloaded files
-        run: du -ah dist
-      - name: Publish package distributions to PyPI
-        if: ${{ !inputs.dry_run }}
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: https://upload.pypi.org/legacy/
@@ -3,12 +3,16 @@ __pycache__/
 *.py[cod]
 *$py.class

+# C extensions
+*.so
+
 # tensor files
 *.safe
 *.safetensors

 # Metal libraries
 *.metallib
+venv/

 # Distribution / packaging
 python/mlx/core
@@ -26,15 +30,18 @@ lib64/
 parts/
 sdist/
 var/
-venv/
 wheels/
 share/python-wheels/
 *.egg-info/
 .installed.cfg
 *.egg
 MANIFEST
-uv.lock
-.DS_Store
+
+# vim
+*.swp
+
+# Ignore build dir
+build/

 # Prerequisites
 *.d
@@ -44,7 +51,6 @@ uv.lock
 *.lo
 *.o
 *.obj
-*.ilk

 # Precompiled Headers
 *.gch
@@ -70,12 +76,9 @@ uv.lock
 *.out
 *.app

-# Debug symbols
-*.pdb
-
-# VSCode
+# VSCode 
 .vscode/
+.DS_Store
+
 # Jetbrains
-.cache/
-# vim
-*.swp
+.cache
@@ -1,22 +1,15 @@
 repos:
-   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v6.0.0
-    hooks:
-    -   id: check-yaml
-    # -   id: end-of-file-fixer
-    # -   id: trailing-whitespace
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v21.1.8
+    rev: v18.1.8
    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: 26.1.0
+    rev: 24.8.0
    hooks:
    -   id: black
-    
 -   repo: https://github.com/pycqa/isort
-    rev: 7.0.0
+    rev: 5.13.2
    hooks:
    -   id: isort
        args:
@@ -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`. Added `cross`. Added `orthogonal` initializer.
+- 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`.
@@ -19,17 +19,11 @@ MLX was developed with contributions from the following individuals:
 - Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
 - Paul Paczuski: Improved stability of BCE loss calculation
 - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops.
- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.

 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
  <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
 </a>

-# Organizations
-
-MLX has received contributions from the following companies:
- NVIDIA Corporation & Affiliates
-
 # Third-Party Software

 MLX leverages several third-party software, listed here together with
@@ -1,32 +1,13 @@
-cmake_minimum_required(VERSION 3.25)
+cmake_minimum_required(VERSION 3.24)

-if(NOT MLX_VERSION)
-  file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
-  string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}")
-  set(_major ${CMAKE_MATCH_1})
-  string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}")
-  set(_minor ${CMAKE_MATCH_1})
-  string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}")
-  set(_patch ${CMAKE_MATCH_1})
-  set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}")
-  set(MLX_VERSION ${MLX_PROJECT_VERSION})
-else()
-  string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION
-                       ${MLX_VERSION})
-endif()
-
-project(
-  mlx
-  LANGUAGES C CXX
-  VERSION ${MLX_PROJECT_VERSION})
+project(mlx LANGUAGES C CXX)

 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
-set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
-set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

 # ----------------------------- Configuration -----------------------------
 option(MLX_BUILD_TESTS "Build tests for mlx" ON)
@@ -35,21 +16,19 @@ option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
 option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
 option(MLX_BUILD_METAL "Build metal backend" ON)
 option(MLX_BUILD_CPU "Build cpu backend" ON)
-option(MLX_BUILD_CUDA "Build cuda backend" OFF)
 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_PYTHON_STUBS "Build stub files for python bindings" ON)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
-option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
-option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF)
-option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF)
-option(USE_UBSAN "Enable UndefinedBehaviorSanitizer (UBSan)" OFF)
-option(USE_TSAN "Enable ThreadSanitizer (TSan)" OFF)
+
+if(NOT MLX_VERSION)
+  set(MLX_VERSION 0.21.0)
+endif()

 # --------------------- Processor tests -------------------------
+
 message(
  STATUS
    "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}"
@@ -70,75 +49,10 @@ if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
      message(WARNING "Building for x86_64 arch is not officially supported.")
    endif()
  endif()
+
 else()
  set(MLX_BUILD_METAL OFF)
-endif()
-
-if(MLX_USE_CCACHE)
-  find_program(CCACHE_PROGRAM ccache)
-  if(CCACHE_PROGRAM)
-    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
-    set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-    set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-  endif()
-endif()
-
-if(USE_ASAN AND USE_TSAN)
-  message(
-    FATAL_ERROR
-      "AddressSanitizer (ASan) and ThreadSanitizer (TSan) are mutually exclusive and cannot be enabled at the same time."
-  )
-endif()
-
-set(SANITIZER_COMPILE_FLAGS "")
-set(SANITIZER_LINK_FLAGS "")
-
-if(USE_ASAN)
-  if(WIN32 AND MSVC)
-    list(APPEND SANITIZER_COMPILE_FLAGS /fsanitize=address)
-    list(APPEND SANITIZER_LINK_FLAGS /fsanitize=address)
-  else()
-    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=address)
-    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=address)
-    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
-      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
-    endif()
-  endif()
-endif()
-
-if(USE_UBSAN)
-  if(WIN32 AND MSVC)
-    if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
-      list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
-      list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
-    else()
-      message(
-        WARNING
-          "UndefinedBehaviorSanitizer (UBSan) is not directly supported via a simple flag in MSVC."
-      )
-    endif()
-  else()
-    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
-    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
-  endif()
-endif()
-
-if(USE_TSAN)
-  if(WIN32 AND MSVC)
-    message(
-      FATAL_ERROR
-        "ThreadSanitizer (TSan) is not supported by the MSVC compiler. Please use Clang or GCC."
-    )
-  elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
-    message(FATAL_ERROR "ThreadSanitizer (TSan) is not supported on macOS.")
-  else()
-    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=thread)
-    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=thread)
-    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
-      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
-    endif()
-  endif()
+  message(WARNING "MLX is prioritised for Apple silicon systems using macOS.")
 endif()

 # ----------------------------- Lib -----------------------------
@@ -149,30 +63,18 @@ cmake_policy(SET CMP0135 NEW)

 add_library(mlx)

-target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
-target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
-
-if(MLX_BUILD_CUDA)
-  enable_language(CUDA)
-  find_package(CUDAToolkit REQUIRED)
-  find_package(CUDNN REQUIRED)
-  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.1" AND CUDAToolkit_VERSION
-                                                          VERSION_LESS "13.2")
-    message(FATAL_ERROR "CUDA Toolkit 13.1 is not supported.")
-  endif()
+if(MLX_BUILD_METAL)
+  set(METAL_LIB "-framework Metal")
+  set(FOUNDATION_LIB "-framework Foundation")
+  set(QUARTZ_LIB "-framework QuartzCore")
 endif()

-if(MLX_BUILD_METAL)
-  find_library(METAL_LIB Metal)
-  find_library(FOUNDATION_LIB Foundation)
-  find_library(QUARTZ_LIB QuartzCore)
-  if(METAL_LIB)
-    message(STATUS "Metal found ${METAL_LIB}")
-  else()
-    message(
-      FATAL_ERROR
-        "Metal not found. Set MLX_BUILD_METAL=OFF to build without GPU")
-  endif()
+if(MLX_BUILD_METAL AND NOT METAL_LIB)
+  message(STATUS "Metal not found. Unable to build GPU")
+  set(MLX_BUILD_METAL OFF)
+  set(MLX_METAL_DEBUG OFF)
+elseif(MLX_BUILD_METAL)
+  message(STATUS "Building METAL sources")

  if(MLX_METAL_DEBUG)
    add_compile_definitions(MLX_METAL_DEBUG)
@@ -181,8 +83,7 @@ if(MLX_BUILD_METAL)
  # Throw an error if xcrun not found
  execute_process(
    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
-    OUTPUT_VARIABLE MACOS_SDK_VERSION
-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
+    OUTPUT_VARIABLE MACOS_SDK_VERSION COMMAND_ERROR_IS_FATAL ANY)

  if(${MACOS_SDK_VERSION} LESS 14.0)
    message(
@@ -192,12 +93,10 @@ if(MLX_BUILD_METAL)
  message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}")

  set(METAL_CPP_URL
-      https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18-beta.zip
+  )

  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
-    if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0)
-      message(FATAL_ERROR "MLX requires macOS >= 14.0")
-    endif()
    set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
  endif()
  execute_process(
@@ -206,6 +105,7 @@ if(MLX_BUILD_METAL)
      "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}>
@@ -213,70 +113,16 @@ if(MLX_BUILD_METAL)
  target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
 endif()

-if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
-  # With newer clang/gcc versions following libs are implicitly linked, but when
-  # building on old distributions they need to be explicitly listed.
-  target_link_libraries(mlx PRIVATE dl pthread)
-endif()
-
-if(WIN32)
-  if(MSVC)
-    # GGUF does not build with MSVC.
-    set(MLX_BUILD_GGUF OFF)
-  endif()
-  # Generate DLL and EXE in the same dir, otherwise EXE will not be able to run.
-  # This is only done when MLX is built as the top project.
-  if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
-    set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
-  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.2
-    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(ACCELERATE_LIBRARY)
    message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
    set(MLX_BUILD_ACCELERATE ON)
-  else()
-    message(STATUS "Accelerate 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(MLX_USE_ACCELERATE)
    add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(WIN32)
-    # Download and link prebuilt binaries of OpenBLAS. Note that we can only
-    # link with the dynamic library, the prebuilt binaries were built with MinGW
-    # so static-linking would require linking with MinGW's runtime.
-    FetchContent_Declare(
-      openblas
-      URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip"
-    )
-    FetchContent_MakeAvailable(openblas)
-    target_link_libraries(mlx
-                          PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib")
-    target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include")
-    # Make sure the DLL file is placed in the same dir with executables.
-    set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll")
-    add_custom_command(
-      TARGET mlx
-      POST_BUILD
-      COMMAND ${CMAKE_COMMAND} -E copy_if_different ${OPENBLAS_DLL_FILE}
-              ${CMAKE_BINARY_DIR})
  else()
+    message(STATUS "Accelerate or arm neon not found, using default backend.")
+    set(MLX_BUILD_ACCELERATE OFF)
    if(${CMAKE_HOST_APPLE})
      # The blas shipped in macOS SDK is not supported, search homebrew for
      # openblas instead.
@@ -294,7 +140,7 @@ if(MLX_BUILD_CPU)
    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 PRIVATE ${LAPACK_LIBRARIES})
+    target_link_libraries(mlx PUBLIC ${LAPACK_LIBRARIES})
    # List blas after lapack otherwise we may accidentally incldue an old
    # version of lapack.h from the include dirs of blas.
    find_package(BLAS REQUIRED)
@@ -307,27 +153,35 @@ if(MLX_BUILD_CPU)
    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 PRIVATE ${BLAS_LIBRARIES})
+    target_link_libraries(mlx PUBLIC ${BLAS_LIBRARIES})
+
+    if(WIN32)
+      find_package(dlfcn-win32 REQUIRED)
+      message(STATUS "dlfcn-win32 lib " ${dlfcn-win32_LIBRARIES})
+      message(STATUS "dlfcn-win32 include " ${dlfcn-win32_INCLUDE_DIRS})
+      target_link_libraries(mlx PUBLIC ${dlfcn-win32_LIBRARIES})
+    endif()
  endif()
 else()
  set(MLX_BUILD_ACCELERATE OFF)
 endif()

-message(STATUS "Downloading json")
-FetchContent_Declare(
-  json
-  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
-FetchContent_MakeAvailable(json)
-target_include_directories(
-  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
-
-# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
-if(MLX_BUILD_CPU
-   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
-   AND DEFINED MACOS_SDK_VERSION
-   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
-  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
-                   ${CMAKE_BINARY_DIR}/jaccl)
+find_package(MPI)
+if(MPI_FOUND)
+  execute_process(
+    COMMAND zsh "-c" "mpirun --version"
+    OUTPUT_VARIABLE MPI_VERSION
+    ERROR_QUIET)
+  if(${MPI_VERSION} MATCHES ".*Open MPI.*")
+    target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
+  elseif(MPI_VERSION STREQUAL "")
+    set(MPI_FOUND FALSE)
+    message(
+      WARNING "MPI found but mpirun is not available. Building without MPI.")
+  else()
+    set(MPI_FOUND FALSE)
+    message(WARNING "MPI which is not OpenMPI found. Building without MPI.")
+  endif()
 endif()

 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
@@ -336,31 +190,26 @@ target_include_directories(
  mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
             $<INSTALL_INTERFACE:include>)

-if(USE_SYSTEM_FMT)
-  find_package(fmt REQUIRED)
-else()
-  FetchContent_Declare(
-    fmt
-    GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-    GIT_TAG 12.1.0
-    EXCLUDE_FROM_ALL)
-  FetchContent_MakeAvailable(fmt)
-endif()
+FetchContent_Declare(
+  fmt
+  GIT_REPOSITORY https://github.com/fmtlib/fmt.git
+  GIT_TAG 10.2.1
+  EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(fmt)
 target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>)

 if(MLX_BUILD_PYTHON_BINDINGS)
  message(STATUS "Building Python bindings.")
  find_package(
-    Python 3.10
+    Python 3.8
    COMPONENTS Interpreter Development.Module
    REQUIRED)
-  FetchContent_Declare(
-    nanobind
-    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v2.12.0
-    GIT_SHALLOW TRUE
-    EXCLUDE_FROM_ALL)
-  FetchContent_MakeAvailable(nanobind)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE NB_DIR)
+  list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
+  find_package(nanobind CONFIG REQUIRED)
  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()

@@ -380,15 +229,6 @@ endif()
 # ----------------------------- Installation -----------------------------
 include(GNUInstallDirs)

-if(WIN32)
-  # Install DLLs to the same dir with extension file (core.pyd) on Windows.
-  set(CMAKE_INSTALL_BINDIR ".")
-  if(MLX_BUILD_CPU)
-    # Install OpenBLAS.
-    install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN)
-  endif()
-endif()
-
 # Install library
 install(
  TARGETS mlx
@@ -5,26 +5,26 @@ possible.

 ## Pull Requests

-1. Fork and submit pull requests to the repo.
+1. Fork and submit pull requests to the repo. 
 2. If you've added code that should be tested, add tests.
 3. If a change is likely to impact efficiency, run some of the benchmarks before
   and after the change. Examples of benchmarks can be found in `benchmarks/python/`.
 4. If you've changed APIs, update the documentation.
-5. Every PR should have passing tests and at least one review.
+5. Every PR should have passing tests and at least one review. 
 6. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`.
   This should install hooks for running `black` and `clang-format` to ensure
   consistent style for C++ and python code.
-
+ 
   You can also run the formatters manually as follows:
-
-   ```shell
-   clang-format -i file.cpp
-   ```
-
-   ```shell
-   black file.py
-   ```
-
+ 
+     ```
+     clang-format -i file.cpp
+     ```
+ 
+     ```
+     black file.py
+     ```
+ 
   or run `pre-commit run --all-files` to check all files in the repo.

 ## Issues
@@ -1,6 +1,4 @@
 include CMakeLists.txt
-include mlx.pc.in
 recursive-include mlx/ *
-include cmake/*
 include python/src/*
 include python/mlx/py.typed # support type hinting as in PEP-561
@@ -2,7 +2,7 @@

 [**Quickstart**](#quickstart) | [**Installation**](#installation) |
 [**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
-[**Examples**](#examples)
+[**Examples**](#examples) 

 [![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)

@@ -11,37 +11,37 @@ brought to you by Apple machine learning research.

 Some key features of MLX include:

- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+ - **Familiar APIs**: MLX has a Python API that closely follows NumPy.  MLX
   also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
   [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
-   the Python API. MLX has higher-level packages like `mlx.nn` and
+   the Python API.  MLX has higher-level packages like `mlx.nn` and
   `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
   more complex models.

- **Composable function transformations**: MLX supports composable function
-  transformations for automatic differentiation, automatic vectorization,
-  and computation graph optimization.
+ - **Composable function transformations**: MLX supports composable function
+   transformations for automatic differentiation, automatic vectorization,
+   and computation graph optimization.

- **Lazy computation**: Computations in MLX are lazy. Arrays are only
-  materialized when needed.
+ - **Lazy computation**: Computations in MLX are lazy. Arrays are only
+   materialized when needed.

- **Dynamic graph construction**: Computation graphs in MLX are constructed
-  dynamically. Changing the shapes of function arguments does not trigger
-  slow compilations, and debugging is simple and intuitive.
+ - **Dynamic graph construction**: Computation graphs in MLX are constructed
+   dynamically. Changing the shapes of function arguments does not trigger
+   slow compilations, and debugging is simple and intuitive.

- **Multi-device**: Operations can run on any of the supported devices
-  (currently the CPU and the GPU).
+ - **Multi-device**: Operations can run on any of the supported devices
+   (currently the CPU and the GPU).

- **Unified memory**: A notable difference from MLX and other frameworks
-  is the *unified memory model*. Arrays in MLX live in shared memory.
-  Operations on MLX arrays can be performed on any of the supported
-  device types without transferring data.
+ - **Unified memory**: A notable difference from MLX and other frameworks
+   is the *unified memory model*. Arrays in MLX live in shared memory.
+   Operations on MLX arrays can be performed on any of the supported
+   device types without transferring data.

 MLX is designed by machine learning researchers for machine learning
 researchers. The framework is intended to be user-friendly, but still efficient
 to train and deploy models. The design of the framework itself is also
 conceptually simple. We intend to make it easy for researchers to extend and
-improve MLX with the goal of quickly exploring new ideas.
+improve MLX with the goal of quickly exploring new ideas. 

 The design of MLX is inspired by frameworks like
 [NumPy](https://numpy.org/doc/stable/index.html),
@@ -68,30 +68,25 @@ in the documentation.

 ## Installation

-MLX is available on [PyPI](https://pypi.org/project/mlx/). To install MLX on
-macOS, run:
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install the Python API, run:

-```bash
+**With `pip`**:
+
+```
 pip install mlx
 ```

-To install the CUDA backend on Linux, run:
+**With `conda`**:

-```bash
-pip install mlx[cuda]
 ```
-
-To install a CPU-only Linux package, run:
-
-```bash
-pip install mlx[cpu]
+conda install -c conda-forge mlx
 ```

 Checkout the
 [documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
 for more information on building the C++ and Python APIs from source.

-## Contributing
+## Contributing 

 Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
 on contributing to MLX. See the
@@ -110,7 +105,7 @@ Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
 MLX useful in your research and wish to cite it, please use the following
 BibTex entry:

-```text
+```
@software{mlx2023,
  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
@@ -5,35 +5,35 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"

-namespace mx = mlx::core;
+using namespace mlx::core;

 void time_value_and_grad() {
-  auto x = mx::ones({200, 1000});
-  mx::eval(x);
-  auto fn = [](mx::array x) {
+  auto x = ones({200, 1000});
+  eval(x);
+  auto fn = [](array x) {
    for (int i = 0; i < 20; ++i) {
-      x = mx::log(mx::exp(x));
+      x = log(exp(x));
    }
-    return mx::sum(x);
+    return sum(x);
  };

-  auto grad_fn = mx::grad(fn);
+  auto grad_fn = grad(fn);
  auto independent_value_and_grad = [&]() {
    auto value = fn(x);
    auto dfdx = grad_fn(x);
-    return std::vector<mx::array>{value, dfdx};
+    return std::vector<array>{value, dfdx};
  };
  TIME(independent_value_and_grad);

-  auto value_and_grad_fn = mx::value_and_grad(fn);
+  auto value_and_grad_fn = value_and_grad(fn);
  auto combined_value_and_grad = [&]() {
    auto [value, dfdx] = value_and_grad_fn(x);
-    return std::vector<mx::array>{value, dfdx};
+    return std::vector<array>{value, dfdx};
  };
  TIME(combined_value_and_grad);
 }

 int main() {
-  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  std::cout << "Benchmarks for " << default_device() << std::endl;
  time_value_and_grad();
 }
@@ -4,21 +4,21 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"

-namespace mx = mlx::core;
+using namespace 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(mx::Device::cpu);
+  set_default_device(Device::cpu);
  for (auto size : sizes) {
-    auto a = mx::random::uniform({size});
-    auto b = mx::random::uniform({size});
-    mx::eval(a, b);
+    auto a = random::uniform({size});
+    auto b = random::uniform({size});
+    eval(a, b);
    std::cout << "Size " << size << std::endl;
-    TIMEM("cpu", mx::add, a, b, mx::Device::cpu);
-    TIMEM("gpu", mx::add, a, b, mx::Device::gpu);
+    TIMEM("cpu", add, a, b, Device::cpu);
+    TIMEM("gpu", add, a, b, Device::gpu);
  }
 }

@@ -1,111 +1,110 @@
 // Copyright © 2023 Apple Inc.

-#include <cstring>
 #include <iostream>
 #include <sstream>

 #include "mlx/mlx.h"
 #include "time_utils.h"

-namespace mx = mlx::core;
+using namespace mlx::core;

 void time_irregular_binary_ops_1D() {
-  auto device = mx::default_device();
+  auto device = default_device();
  int size = 1000000;
  int step = 2;
-  auto a = mx::random::uniform({size});
-  auto b = mx::random::uniform({size});
-  mx::eval(a, b);
+  auto a = random::uniform({size});
+  auto b = random::uniform({size});
+  eval(a, b);
  a = slice(a, {0}, {size}, {step});
  b = slice(b, {0}, {size}, {step});
-  TIMEM("1D strided", mx::add, a, b, device);
+  TIMEM("1D strided", add, a, b, device);
 }

 void time_irregular_binary_ops_2D() {
-  auto device = mx::default_device();
+  auto device = default_device();
  int size = 2048;
-  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);
+  auto a = random::uniform({size, size});
+  auto b = random::uniform({size, size});
+  eval(a, b);
+  TIMEM("2D regular", add, a, b, device);

-  b = mx::transpose(b);
-  mx::eval(b);
-  TIMEM("2D mx::transpose", mx::add, a, b, device);
+  b = transpose(b);
+  eval(b);
+  TIMEM("2D transpose", add, a, b, device);

-  b = mx::random::uniform({size});
-  mx::eval(b);
-  TIMEM("2D broadcast dim 0", mx::add, a, b, device);
+  b = random::uniform({size});
+  eval(b);
+  TIMEM("2D broadcast dim 0", add, a, b, device);

-  b = mx::reshape(b, {size, 1});
-  mx::eval(b);
-  TIMEM("2D broadcast dim 1", mx::add, a, b, device);
+  b = reshape(b, {size, 1});
+  eval(b);
+  TIMEM("2D broadcast dim 1", add, a, b, device);
 }

 void time_irregular_binary_ops_3D() {
-  auto device = mx::default_device();
+  auto device = default_device();
  int d0 = 32;
  int d1 = 512;
  int d2 = 512;
-  auto a = mx::random::uniform({d0, d1, d2});
-  auto b = mx::random::uniform({d0, d1, d2});
-  TIMEM("3D regular", mx::add, a, b, device);
+  auto a = random::uniform({d0, d1, d2});
+  auto b = random::uniform({d0, d1, d2});
+  TIMEM("3D regular", add, a, b, device);

-  b = mx::transpose(b, {0, 2, 1});
-  TIMEM("3D mx::transpose", mx::add, a, b, device);
+  b = transpose(b, {0, 2, 1});
+  TIMEM("3D transpose", add, a, b, device);

-  b = mx::random::uniform({d1, d2});
-  TIMEM("3D broadcast dim 0", mx::add, a, b, device);
+  b = random::uniform({d1, d2});
+  TIMEM("3D broadcast dim 0", 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, 1, d2});
+  TIMEM("3D broadcast dim 1", add, a, b, device);

-  b = mx::random::uniform({d0, d1, 1});
-  TIMEM("3D broadcast dim 2", mx::add, a, b, device);
+  b = random::uniform({d0, d1, 1});
+  TIMEM("3D broadcast dim 2", add, a, b, device);

-  b = mx::random::uniform({d2});
-  TIMEM("3D broadcast dims 0, 1", mx::add, a, b, device);
+  b = random::uniform({d2});
+  TIMEM("3D broadcast dims 0, 1", add, a, b, device);

-  b = mx::random::uniform({d1, 1});
-  TIMEM("3D broadcast dims 0, 2", mx::add, a, b, device);
+  b = random::uniform({d1, 1});
+  TIMEM("3D broadcast dims 0, 2", add, a, b, device);

-  b = mx::random::uniform({d0, 1, 1});
-  TIMEM("3D broadcast dims 1, 2", mx::add, a, b, device);
+  b = random::uniform({d0, 1, 1});
+  TIMEM("3D broadcast dims 1, 2", add, a, b, device);
 }

 void time_irregular_binary_ops_4D() {
-  auto device = mx::default_device();
-  mx::Shape shape = {8, 8, 512, 512};
-  auto a = mx::random::uniform(shape);
-  auto b = mx::random::uniform(shape);
+  auto device = default_device();
+  std::vector<int> shape = {8, 8, 512, 512};
+  auto a = random::uniform(shape);
+  auto b = random::uniform(shape);

-  TIMEM("4D regular", mx::add, a, b, device);
+  TIMEM("4D regular", add, a, b, device);

-  b = mx::transpose(b, {0, 1, 3, 2});
-  TIMEM("4D mx::transpose", mx::add, a, b, device);
+  b = transpose(b, {0, 1, 3, 2});
+  TIMEM("4D transpose", add, a, b, device);

  std::string om = "4D broadcast dims ";
  for (int i = 0; i < shape.size(); ++i) {
    shape[i] = 1;
-    b = mx::random::uniform(shape);
+    b = random::uniform(shape);
    std::ostringstream msg;
    msg << om << i;
-    TIMEM(msg.str(), mx::add, a, b, device);
+    TIMEM(msg.str(), add, a, b, device);

    for (int j = i + 1; j < shape.size(); ++j) {
      shape[j] = 1;
      std::ostringstream msg;
      msg << om << i << ", " << j;
-      b = mx::random::uniform(shape);
-      TIMEM(msg.str(), mx::add, a, b, device);
+      b = random::uniform(shape);
+      TIMEM(msg.str(), 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 = mx::random::uniform(shape);
-        TIMEM(msg.str(), mx::add, a, b, device);
+        b = random::uniform(shape);
+        TIMEM(msg.str(), add, a, b, device);
        shape[k] = a.shape(k);
      }
    }
@@ -114,83 +113,83 @@ void time_irregular_binary_ops_4D() {
 }

 void time_irregular_reshape() {
-  auto device = mx::default_device();
-  mx::Shape shape;
-  auto reshape_fn = [&shape, device](const mx::array& a) {
-    return mx::reshape(a, shape, device);
+  auto device = default_device();
+  std::vector<int> shape;
+  auto reshape_fn = [&shape, device](const array& a) {
+    return reshape(a, shape, device);
  };

  int size = 64;
  int d = 2 * size;

-  auto a = mx::random::uniform({d, d, d});
+  auto a = random::uniform({d, d, d});

  shape = {8 * size, size, size};
  TIMEM("3D contiguous", reshape_fn, a);

-  a = mx::transpose(a);
+  a = transpose(a);
  shape = {8 * size, size, size};
-  TIMEM("3D mx::transpose", reshape_fn, a);
+  TIMEM("3D transpose", reshape_fn, a);

-  a = mx::transpose(a, {1, 2, 0});
+  a = transpose(a, {1, 2, 0});
  shape = {8 * size, size, size};
-  TIMEM("3D mx::transpose dims 1 2", reshape_fn, a);
+  TIMEM("3D transpose dims 1 2", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d, d}), {d, d, d});
+  a = broadcast_to(random::uniform({d, d}), {d, d, d});
  TIMEM("3D broadcast dim 0", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d, 1, d}), {d, d, d});
+  a = broadcast_to(random::uniform({d, 1, d}), {d, d, d});
  TIMEM("3D broadcast dim 1", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d, d, 1}), {d, d, d});
+  a = broadcast_to(random::uniform({d, d, 1}), {d, d, d});
  TIMEM("3D broadcast dim 2", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d}), {d, d, d});
+  a = broadcast_to(random::uniform({d}), {d, d, d});
  TIMEM("3D broadcast dims 0, 1", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d, 1}), {d, d, d});
+  a = broadcast_to(random::uniform({d, 1}), {d, d, d});
  TIMEM("3D broadcast dims 0, 2", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({d, 1, 1}), {d, d, d});
+  a = broadcast_to(random::uniform({d, 1, 1}), {d, d, d});
  TIMEM("3D broadcast dims 1, 2", reshape_fn, a);

-  a = mx::broadcast_to(mx::random::uniform({1, 1, 1}), {d, d, d});
+  a = broadcast_to(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 = mx::default_device();
+  auto device = default_device();
  int size = 1000000;
  int step = 2;
-  auto a = mx::random::uniform({size});
+  auto a = random::uniform({size});
  a = slice(a, {0}, {size}, {step});
-  TIMEM("1D strided", mx::astype, a, mx::int32, device);
+  TIMEM("1D strided", astype, a, int32, device);
 }

 void time_irregular_astype_2D() {
-  auto device = mx::default_device();
+  auto device = default_device();
  int size = 2048;
-  mx::Shape shape = {size, size};
+  std::vector<int> shape = {size, size};

-  auto a = mx::random::uniform(shape);
-  TIMEM("2D regular", mx::astype, a, mx::int32, device);
+  auto a = random::uniform(shape);
+  TIMEM("2D regular", astype, a, int32, device);

-  a = mx::transpose(a);
-  TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device);
+  a = transpose(a);
+  TIMEM("2D transpose", 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}), shape);
+  TIMEM("2D broadcast dim 0", 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);
+  a = broadcast_to(random::uniform({size, 1}), shape);
+  TIMEM("2D broadcast dim 1", astype, a, int32, device);
 }

 int main(int argc, char** argv) {
  if (argc > 1) {
    bool use_gpu = !strcmp(argv[1], "gpu");
-    set_default_device(use_gpu ? mx::Device::gpu : mx::Device::cpu);
+    set_default_device(use_gpu ? Device::gpu : Device::cpu);
  }
-  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  std::cout << "Benchmarks for " << default_device() << std::endl;
  time_irregular_binary_ops_1D();
  time_irregular_binary_ops_2D();
  time_irregular_binary_ops_3D();
@@ -3,20 +3,20 @@
 #include "mlx/mlx.h"
 #include "time_utils.h"

-namespace mx = mlx::core;
+using namespace mlx::core;

 void time_creation_ops() {
  int M = 2000;
  int N = 500;
  auto shape = {M, N};
-  auto full_fp32 = [&]() { return mx::full(shape, 3.3f); };
+  auto full_fp32 = [&]() { return full(shape, 3.3f); };
  TIME(full_fp32);
-  auto zeros_fp32 = [&]() { return mx::zeros(shape, mx::float32); };
+  auto zeros_fp32 = [&]() { return zeros(shape, float32); };
  TIME(zeros_fp32);
-  auto ones_fp32 = [&]() { return mx::ones(shape, mx::float32); };
+  auto ones_fp32 = [&]() { return ones(shape, float32); };
  TIME(ones_fp32);

-  auto arange_fp32 = [&]() { return mx::arange(0.0, 10.0, 1e-4); };
+  auto arange_fp32 = [&]() { return arange(0.0, 10.0, 1e-4); };
  TIME(arange_fp32);
 }

@@ -24,212 +24,194 @@ void time_type_conversions() {
  int M = 2000;
  int N = 500;
  auto shape = {M, N};
-  auto device = mx::default_device();
+  auto device = default_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);
+  auto a = zeros(shape, float32);
+  eval(a);
+  TIMEM("float32 to int32", astype, a, int32, device);
+  TIMEM("float32 to uint32", astype, a, uint32, 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, int32);
+  eval(a);
+  TIMEM("int32 to float32", astype, a, float32, 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);
+  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);
 }

 void time_random_generation() {
  int M = 2000;
  int N = 500;

-  auto uniform = [&]() { return mx::random::uniform({M, N}, mx::float32); };
+  auto uniform = [&]() { return random::uniform({M, N}, float32); };
  TIME(uniform);
-  auto normal = [&]() { return mx::random::normal({M, N}, mx::float32); };
+  auto normal = [&]() { return random::normal({M, N}, float32); };
  TIME(normal);
 }

 void time_unary_ops() {
  int M = 2000;
  int N = 500;
-  auto device = mx::default_device();
+  auto device = default_device();

-  auto a = mx::random::normal({M, N});
-  mx::eval(a);
+  auto a = random::normal({M, N});
+  eval(a);
  TIME(mlx::core::abs, a, device);
-  TIME(mx::negative, a, device);
-  TIME(mx::sign, a, device);
-  TIME(mx::square, a, device);
+  TIME(negative, a, device);
+  TIME(sign, a, device);
+  TIME(square, a, device);
  TIME(mlx::core::sqrt, a, device);
-  TIME(mx::rsqrt, a, device);
+  TIME(rsqrt, a, device);
  TIME(mlx::core::exp, a, device);

-  a = mx::random::uniform({M, N});
+  a = random::uniform({M, N});
  TIME(mlx::core::log, a, device);
 }

 void time_binary_ops() {
  int M = 1000, N = 100, K = 10;
-  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);
+  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);

-  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);
+  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);

-  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 = 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::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);
+  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);
 }

 void time_strided_ops() {
  int M = 50, N = 50, O = 50, P = 50;
-  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);
+  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);
 }

 void time_comparisons() {
  int M = 1000, N = 100, K = 10;
-  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);
+  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);
 }

 void time_matvec() {
  int M = 2000, N = 200;
-  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); };
+  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); };
  TIME(matvec);

-  auto matvec_transpose = [&]() { return mx::matmul(mx::transpose(a), c); };
+  auto matvec_transpose = [&]() { return matmul(transpose(a), c); };
  TIME(matvec_transpose);
 }

 void time_matmul() {
  int M = 1000, N = 1000, K = 1000;
-  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 a = random::uniform({M, K});
+  auto b = random::uniform({K, N});
+  auto device = default_device();
+  eval(a, b);
+  TIME(matmul, a, b, device);

-  auto transpose_matmul = [&]() { return mx::matmul(mx::transpose(a), b); };
+  auto transpose_matmul = [&]() { return matmul(transpose(a), b); };
  TIME(transpose_matmul);
 }

 void time_reductions() {
-  auto a = mx::random::normal({10000, 1000});
-  mx::eval(a);
-  auto sum_all = [&a]() { return mx::sum(a, false); };
+  auto a = random::normal({10000, 1000});
+  eval(a);
+  auto sum_all = [&a]() { return sum(a, false); };
  TIME(sum_all);

-  auto sum_along_0 = [&a]() { return mx::sum(a, 0, false); };
+  auto sum_along_0 = [&a]() { return sum(a, 0, false); };
  TIME(sum_along_0);

-  auto sum_along_1 = [&a]() { return mx::sum(a, 1, false); };
+  auto sum_along_1 = [&a]() { return sum(a, 1, false); };
  TIME(sum_along_1);

-  auto prod_all = [&a]() { return mx::prod(a, false); };
+  auto prod_all = [&a]() { return prod(a, false); };
  TIME(prod_all);

-  auto all_true = [&a]() { return mx::all(a, false); };
+  auto all_true = [&a]() { return all(a, false); };
  TIME(all_true);

-  auto all_along_0 = [&a]() { return mx::all(a, 0, false); };
+  auto all_along_0 = [&a]() { return all(a, 0, false); };
  TIME(all_along_0);

-  auto all_along_1 = [&a]() { return mx::all(a, 1, false); };
+  auto all_along_1 = [&a]() { return all(a, 1, false); };
  TIME(all_along_1);

-  auto any_true = [&a]() { return mx::any(a, false); };
+  auto any_true = [&a]() { return any(a, false); };
  TIME(any_true);

-  auto argmin_along_0 = [&a]() { return mx::argmin(a, 0, false); };
+  auto argmin_along_0 = [&a]() { return argmin(a, 0, false); };
  TIME(argmin_along_0);

-  auto argmin_along_1 = [&a]() { return mx::argmin(a, 1, false); };
+  auto argmin_along_1 = [&a]() { return argmin(a, 1, false); };
  TIME(argmin_along_1);
-
-  auto indices = mx::array({1});
-  auto updates = mx::reshape(mx::array({NAN}), {1, 1, 1});
-  std::vector<int> axes{0};
-  auto b = scatter(a, {indices}, updates, axes);
-  mx::eval(b);
-
-  auto max_along_0 = [&b]() { return mx::max(b, 0, false); };
-  TIME(max_along_0);
-  auto max_along_1 = [&b]() { return mx::max(b, 1, false); };
-  TIME(max_along_1);
-
-  auto min_along_0 = [&b]() { return mx::min(b, 0, false); };
-  TIME(min_along_0);
-  auto min_along_1 = [&b]() { return mx::min(b, 1, false); };
-  TIME(min_along_1);
 }

 void time_gather_scatter() {
-  auto a = mx::random::normal({1000, 768});
-  mx::eval(a);
-  auto indices = mx::random::randint(0, 1000, {256});
-  mx::eval(indices);
+  auto a = random::normal({1000, 768});
+  eval(a);
+  auto indices = random::randint(0, 1000, {256});
+  eval(indices);

-  auto embedding_lookup = [&a, &indices]() { return mx::take(a, indices, 0); };
+  auto embedding_lookup = [&a, &indices]() { return take(a, indices, 0); };
  TIME(embedding_lookup);

-  indices = mx::random::randint(0, 768 * 1000, {256 * 768});
-  mx::eval(indices);
+  indices = random::randint(0, 768 * 1000, {256 * 768});
+  eval(indices);

-  auto single_element_lookup = [&a, &indices]() {
-    return mx::take(a, indices);
-  };
+  auto single_element_lookup = [&a, &indices]() { return take(a, indices); };
  TIME(single_element_lookup);

-  indices = mx::random::randint(0, 1000, {256});
-  auto updates = mx::random::normal({256, 1, 768});
-  mx::eval(indices, updates);
+  indices = random::randint(0, 1000, {256});
+  auto updates = random::normal({256, 1, 768});
+  eval(indices, updates);

  auto embedding_update = [&a, &indices, &updates]() {
    return scatter(a, indices, updates, 0);
@@ -241,10 +223,10 @@ void time_gather_scatter() {
  };
  TIME(embedding_add);

-  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);
+  a = reshape(a, {-1});
+  indices = random::randint(0, 768 * 1000, {768 * 256});
+  updates = random::normal({256 * 768, 1});
+  eval(a, indices, updates);

  auto single_element_update = [&a, &indices, &updates]() {
    return scatter(a, indices, updates, 0);
@@ -258,21 +240,21 @@ void time_gather_scatter() {
 }

 void time_divmod() {
-  auto a = mx::random::normal({1000});
-  auto b = mx::random::normal({1000});
-  mx::eval({a, b});
+  auto a = random::normal({1000});
+  auto b = random::normal({1000});
+  eval({a, b});

-  auto divmod_fused = [&a, &b]() { return mx::divmod(a, b); };
+  auto divmod_fused = [&a, &b]() { return divmod(a, b); };
  TIME(divmod_fused);

  auto divmod_separate = [&a, &b]() {
-    return std::vector<mx::array>{mx::floor_divide(a, b), mx::remainder(a, b)};
+    return std::vector<array>{floor_divide(a, b), remainder(a, b)};
  };
  TIME(divmod_separate);
 }

 int main() {
-  std::cout << "Benchmarks for " << mx::default_device() << std::endl;
+  std::cout << "Benchmarks for " << default_device() << std::endl;
  time_creation_ops();
  time_type_conversions();
  time_unary_ops();
@@ -142,7 +142,9 @@ def bench_shape(B, M, N, K, np_dtype, transpose="nn"):
    t_b = (0, 1, 2) if transpose[1] == "n" else (0, 2, 1)

    c_mlx = a_mx.transpose(t_a) @ b_mx.transpose(t_b)
-    c_npy = a_np.transpose(t_a).astype(np_dtype) @ b_np.transpose(t_b).astype(np_dtype)
+    c_npy = a_np.transpose(t_a).astype(np.float32) @ b_np.transpose(t_b).astype(
+        np.float32
+    )

    atol = 1e-5 if np_dtype == np.float32 else 1e-4

@@ -161,7 +163,7 @@ def get_gflop_count(B, M, N, K):
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run gemm benchmarks")

-    dtypes = ("float32", "float16", "complex64")
+    dtypes = ("float32", "float16")
    transposes = ("nn", "nt", "tn")
    shapes = (
        (16, 234, 768, 3072),
@@ -185,7 +187,7 @@ if __name__ == "__main__":
                diff = gflops_mx / gflops_pt - 1.0

                print(
-                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%"
+                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100. * diff:+5.2f}%"
                )
                if gflops_pt >= 2.0 * gflops_mx:
                    print("ATTENTION ^^^^^^^")
@@ -1,5 +1,6 @@
 # Copyright © 2023 Apple Inc.

+import argparse
 import os
 import subprocess
 import time
@@ -195,7 +196,7 @@ def bench_with_out_len(ax, out_vec_len, in_vector_lens, dtype, transpose):


 for transpose in (False, True):
-    for dtype in ("float32", "float16", "complex64"):
+    for dtype in ("float32", "float16"):
        fig, axs = plt.subplots(
            len(in_vec_sizes), 2, figsize=(8.5, 11), layout="constrained"
        )
@@ -214,7 +215,7 @@ for transpose in (False, True):
        fig.suptitle(f"{device_name}: {dtype} {op_name}")
        fig.savefig(
            os.path.join(
-                results_dir, f"{device_name.replace(' ', '_')}_{dtype}_{op_name}.pdf"
+                results_dir, f'{device_name.replace(" ", "_")}_{dtype}_{op_name}.pdf'
            )
        )
        plt.close(fig)
@@ -1,193 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-import time
-
-import mlx.core as mx
-import numpy as np
-
-MLX_DTYPES = {
-    "float16": mx.float16,
-    "bfloat16": mx.bfloat16,
-    "float32": mx.float32,
-}
-
-
-def parse_cases(cases):
-    parsed = []
-    for spec in cases.split(","):
-        parts = spec.split("x")
-        m, n, k, bs = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
-        sparsity = float(parts[4]) if len(parts) > 4 else 0.5
-        parsed.append((m, n, k, bs, sparsity))
-    return parsed
-
-
-def make_masks(m, n, k, block_size, sparsity, rng):
-    """Create block masks with given sparsity (fraction of blocks zeroed)."""
-    tm = (m + block_size - 1) // block_size
-    tn = (n + block_size - 1) // block_size
-    tk = (k + block_size - 1) // block_size
-
-    lhs_mask = (rng.random((tm, tk)) >= sparsity).astype(np.bool_)
-    rhs_mask = (rng.random((tk, tn)) >= sparsity).astype(np.bool_)
-    out_mask = (rng.random((tm, tn)) >= sparsity).astype(np.bool_)
-    return lhs_mask, rhs_mask, out_mask
-
-
-def mlx_naive_block_masked_mm(a, b, block_size, out_mask, lhs_mask, rhs_mask):
-    """MLX naive: expand masks and use regular matmul."""
-    M, K = a.shape[-2], a.shape[-1]
-    N = b.shape[-1]
-
-    def expand(mask, rows, cols):
-        e = mx.repeat(mx.repeat(mask, block_size, axis=-2), block_size, axis=-1)
-        return e[..., :rows, :cols]
-
-    a_masked = a * expand(lhs_mask, M, K)
-    b_masked = b * expand(rhs_mask, K, N)
-    c = a_masked @ b_masked
-    c = c * expand(out_mask, M, N)
-    return c
-
-
-def bench_mlx(fn, warmup, iters):
-    for _ in range(warmup):
-        y = fn()
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = fn()
-        mx.eval(y)
-    mx.synchronize()
-    return (time.perf_counter() - start) * 1e3 / iters
-
-
-def print_table(headers, rows):
-    widths = [len(h) for h in headers]
-    for row in rows:
-        for i, cell in enumerate(row):
-            widths[i] = max(widths[i], len(cell))
-
-    def fmt_row(row):
-        return (
-            "| "
-            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
-            + " |"
-        )
-
-    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
-    print(fmt_row(headers))
-    print(sep)
-    for row in rows:
-        print(fmt_row(row))
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Benchmark block_masked_mm vs naive expand+matmul"
-    )
-    parser.add_argument(
-        "--cases",
-        default=(
-            "256x256x256x32x0.5,"
-            "512x512x512x32x0.5,"
-            "1024x1024x1024x32x0.5,"
-            "1024x1024x1024x64x0.5,"
-            "2048x2048x2048x64x0.5,"
-            "256x256x256x32x0.0,"
-            "1024x1024x1024x32x0.0,"
-            "1024x1024x1024x32x0.9"
-        ),
-        help="Comma-separated MxNxKxBSxSparsity list. Sparsity=fraction of blocks zeroed.",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="float32",
-        choices=["float16", "bfloat16", "float32"],
-    )
-    parser.add_argument("--warmup", type=int, default=10)
-    parser.add_argument("--iters", type=int, default=50)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--no-check", action="store_true")
-    args = parser.parse_args()
-
-    mlx_dtype = MLX_DTYPES[args.dtype]
-
-    print(f"dtype={args.dtype}  warmup={args.warmup}  iters={args.iters}")
-
-    headers = [
-        "Case (MxNxKxBS)",
-        "Sparsity",
-        "MLX ms",
-        "Naive ms",
-        "Speedup",
-    ]
-    if not args.no_check:
-        headers.append("Max err")
-    rows = []
-
-    cases = parse_cases(args.cases)
-    for idx, (m, n, k, bs, sparsity) in enumerate(cases):
-        rng = np.random.default_rng(args.seed + idx)
-        a_np = rng.standard_normal((m, k)).astype(np.float32)
-        b_np = rng.standard_normal((k, n)).astype(np.float32)
-        lhs_mask_np, rhs_mask_np, out_mask_np = make_masks(m, n, k, bs, sparsity, rng)
-
-        a_mx = mx.array(a_np, dtype=mlx_dtype)
-        b_mx = mx.array(b_np, dtype=mlx_dtype)
-        lhs_mask_mx = mx.array(lhs_mask_np)
-        rhs_mask_mx = mx.array(rhs_mask_np)
-        out_mask_mx = mx.array(out_mask_np)
-        mx.eval(a_mx, b_mx, lhs_mask_mx, rhs_mask_mx, out_mask_mx)
-
-        # Correctness check: block_masked_mm vs naive expand+matmul
-        err_str = ""
-        if not args.no_check:
-            y_op = mx.block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            )
-            y_naive = mlx_naive_block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            )
-            mx.eval(y_op, y_naive)
-            err = float(mx.max(mx.abs(y_op - y_naive)).item())
-            err_str = f"{err:.2e}"
-
-        # Benchmark
-        t_mlx = bench_mlx(
-            lambda: mx.block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            ),
-            args.warmup,
-            args.iters,
-        )
-        t_naive = bench_mlx(
-            lambda: mlx_naive_block_masked_mm(
-                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
-            ),
-            args.warmup,
-            args.iters,
-        )
-        speedup = f"{t_naive / t_mlx:.2f}x" if t_mlx > 0 else "-"
-
-        row = [
-            f"{m}x{n}x{k}x{bs}",
-            f"{sparsity:.0%}",
-            f"{t_mlx:.3f}",
-            f"{t_naive:.3f}",
-            speedup,
-        ]
-        if not args.no_check:
-            row.append(err_str)
-        rows.append(row)
-
-    print_table(headers, rows)
-    if not args.no_check:
-        print("err: max|block_masked_mm - naive_expand_matmul|")
-
-
-if __name__ == "__main__":
-    main()
@@ -38,10 +38,10 @@ def bench(f, *args):
    for i in range(10):
        f(*args)

-    s = time.perf_counter()
+    s = time.time()
    for i in range(100):
        f(*args)
-    e = time.perf_counter()
+    e = time.time()
    return e - s


@@ -5,7 +5,6 @@ import os
 import time

 import torch
-import torch.cuda
 import torch.mps


@@ -37,18 +36,16 @@ def bench(f, *args):
    for i in range(10):
        f(*args)

-    s = time.perf_counter()
+    s = time.time()
    for i in range(100):
        f(*args)
-    e = time.perf_counter()
+    e = time.time()
    return e - s


 def sync_if_needed(x):
-    if x.device == torch.device("mps"):
+    if x.device != torch.device("cpu"):
        torch.mps.synchronize()
-    elif x.device == torch.device("cuda"):
-        torch.cuda.synchronize()


@torch.no_grad()
@@ -102,14 +99,6 @@ def reduction(op, axis, x):
    sync_if_needed(x)


-@torch.no_grad()
-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)
-    sync_if_needed(x)
-
-
@torch.no_grad()
 def softmax(axis, x):
    ys = []
@@ -351,11 +340,7 @@ if __name__ == "__main__":
        args.axis.pop(0)

    torch.set_num_threads(1)
-    device = "mps"
-    if torch.cuda.is_available():
-        device = "cuda"
-    if args.cpu:
-        device = "cpu"
+    device = "cpu" if args.cpu else "mps"

    types = args.dtype
    if not types:
@@ -475,8 +460,5 @@ 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(f"Unknown benchmark `{args.benchmark}`.")
@@ -1,152 +0,0 @@
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_warmup = 2
-N_iter_bench = 10
-N_iter_func = 10
-
-
-def bench(f, a, b, b_prime):
-    for i in range(N_warmup):
-        f(a, b, b_prime)
-    torch.mps.synchronize()
-
-    s = time.perf_counter_ns()
-    for i in range(N_iter_bench):
-        f(a, b, b_prime)
-    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, b_prime):
-        y = a
-        for i in range(N_iter_func):
-            y = mx.conv3d(y, b, stride=strides, padding=padding, groups=groups)
-            y = mx.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
-        mx.eval(y)
-        return y
-
-    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, b_prime):
-        y = a
-        for i in range(N_iter_func):
-            y = torch.conv3d(y, b, stride=strides, padding=padding, groups=groups)
-            y = torch.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
-        torch.mps.synchronize()
-        return y
-
-    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))
-    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups)))
-    b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(O / groups)))
-
-    a_np, b_np, b_prime_np = map(lambda x: x.astype(np_dtype), (a_np, b_np, b_prime_np))
-    a_mx, b_mx, b_prime_mx = map(lambda x: mx.array(x), (a_np, b_np, b_prime_np))
-    a_pt, b_pt, b_prime_pt = map(
-        lambda x: torch.from_numpy(x.transpose(0, 4, 1, 2, 3)).to("mps"),
-        (a_np, b_np, b_prime_np),
-    )
-
-    torch.mps.synchronize()
-
-    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, b_prime_pt)
-    time_mlx = bench(f_mx, a_mx, b_mx, b_prime_mx)
-
-    # Measure MLX memory
-    mx.clear_cache()
-    mx.reset_peak_memory()
-    y = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
-    mx.eval(y)
-    mlx_peak_mb = mx.get_peak_memory() / 1024**2
-    mlx_active_mb = mx.get_active_memory() / 1024**2
-    del y
-
-    # Measure PyTorch MPS memory
-    torch.mps.synchronize()
-    torch.mps.empty_cache()
-    y = torch.conv3d(a_pt, b_pt, stride=strides, padding=padding, groups=groups)
-    torch.mps.synchronize()
-    pt_current_mb = torch.mps.current_allocated_memory() / 1024**2
-    pt_driver_mb = torch.mps.driver_allocated_memory() / 1024**2
-    del y
-
-    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 5e-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)} "
-            f"[strides = {strides}, padding = {padding}, groups = {groups}] "
-            f"with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
-        )
-
-    return time_mlx, time_torch, mlx_peak_mb, mlx_active_mb, pt_current_mb, pt_driver_mb
-
-
-if __name__ == "__main__":
-    dtypes = ("float16", "float32")
-    shapes = (
-        # (C % 16 == 0)
-        (4, 16, 16, 16, 32, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        # Larger spatial dims
-        (2, 64, 64, 64, 32, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
-        (1, 64, 64, 64, 64, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
-        # Strided
-        (4, 32, 32, 32, 64, 3, 3, 3, 128, (2, 2, 2), (1, 1, 1), 1),
-        # Asymmetric kernels
-        (4, 32, 32, 32, 64, 3, 1, 1, 128, (1, 1, 1), (1, 0, 0), 1),
-        (4, 32, 32, 32, 64, 1, 3, 3, 128, (1, 1, 1), (0, 1, 1), 1),
-        # (C % 16 != 0)
-        (4, 16, 16, 16, 21, 3, 3, 3, 21, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
-        (4, 32, 32, 32, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
-        (4, 16, 16, 16, 3, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
-    )
-
-    for dtype in dtypes:
-        print(f"\n{'=' * 120}" f"\n  dtype: {dtype}" f"\n{'=' * 120}")
-        print(
-            f"{'(N,   D,   H,   W,   C)':<26s} {'(  O, kD, kH, kW,   C)':<24s} "
-            f"{'stride':<12s} {'pads':<12s} {'groups':>6s} "
-            f"{'diff%':>7s}  "
-            f"{'MLX peak':>9s} {'MLX act':>8s} {'PT cur':>8s} {'PT drv':>8s}"
-        )
-        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
-            np_dtype = getattr(np, dtype)
-            time_mlx, time_torch, mlx_peak, mlx_act, pt_cur, pt_drv = 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}), "
-                f"{strides}, {padding}, {groups:6d}, "
-                f"{100. * diff:+6.1f}%  "
-                f"{mlx_peak:8.1f}  {mlx_act:7.1f}  {pt_cur:7.1f}  {pt_drv:7.1f}"
-            )
@@ -1,107 +0,0 @@
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_warmup = 10
-N_iter_bench = 100
-N_iter_func = 5
-
-
-def bench(f, a, b):
-    for i in range(N_warmup):
-        f(a, b)
-    torch.mps.synchronize()
-
-    s = time.perf_counter_ns()
-    for i in range(N_iter_bench):
-        f(a, b)
-    e = time.perf_counter_ns()
-    return (e - s) * 1e-9
-
-
-def make_mx_conv_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)
-        torch.mps.synchronize()
-        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("mps")
-    b_pt = torch.from_numpy(b_np.transpose((0, 3, 1, 2))).to("mps")
-
-    torch.mps.synchronize()
-
-    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__":
-    dtype = "float32"
-    shapes = (
-        (4, 32, 32, 21, 3, 3, 128),
-        (4, 32, 32, 21, 3, 3, 37),
-        (4, 32, 32, 370, 3, 3, 370),
-        (4, 32, 32, 370, 7, 7, 128),
-        (2, 320, 640, 21, 7, 7, 21),
-    )
-    for N, H, W, C, kh, kw, O in shapes:
-        time_mlx, time_torch = bench_shape(
-            N, H, W, C, kh, kw, O, (1, 1), (0, 0), 1, 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}, {100. * diff:+5.2f}%"
-        )
-        if time_mlx >= 2.0 * time_torch:
-            print("ATTENTION ^^^^^^^")
@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 import argparse
+from time import time

 import mlx.core as mx
 import torch
@@ -1,74 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate([x[i] @ w[j].T for i, j in enumerate(idx.tolist())], axis=0)
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_mm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_mm(x, w1.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_mm(x, w2.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = x @ w1.T
-        x = x @ w2.T
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_mm()
@@ -1,84 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate(
-            [
-                mx.quantized_matmul(x[i], w[0][j], w[1][j], w[2][j], transpose=True)
-                for i, j in enumerate(idx.tolist())
-            ],
-            axis=0,
-        )
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_qmm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = mx.quantized_matmul(x, *w1, transpose=True)
-        x = mx.quantized_matmul(x, *w2, transpose=True)
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_qmm()
@@ -1,119 +0,0 @@
-# Copyright © 2026 Apple Inc.
-
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_WARMUP = 5
-N_BENCH = 20
-
-
-def bench_mlx(a, b):
-    for _ in range(N_WARMUP):
-        mx.eval(a @ b)
-
-    times = []
-    for _ in range(N_BENCH):
-        start = time.perf_counter_ns()
-        mx.eval(a @ b)
-        end = time.perf_counter_ns()
-        times.append((end - start) * 1e-9)
-
-    return np.mean(times), np.std(times)
-
-
-@torch.no_grad()
-def bench_torch(a, b):
-    for _ in range(N_WARMUP):
-        _ = a @ b
-        torch.mps.synchronize()
-
-    times = []
-    for _ in range(N_BENCH):
-        start = time.perf_counter_ns()
-        _ = a @ b
-        torch.mps.synchronize()
-        end = time.perf_counter_ns()
-        times.append((end - start) * 1e-9)
-
-    return np.mean(times), np.std(times)
-
-
-def check_correctness(out_mx, out_pt, rtol, M, N, K):
-    if not np.allclose(out_pt, out_mx, rtol=rtol, atol=0):
-        abs_diff = np.abs(out_pt - out_mx)
-        rel_diff = abs_diff / np.maximum(np.abs(out_pt), 1e-10)
-
-        print(
-            f"  WARNING: Correctness failed at {M}x{N}x{K}: "
-            f"max_abs={np.max(abs_diff):.6e}, max_rel={np.max(rel_diff):.6e}"
-        )
-
-
-def bench_gemm(M, N, K, dtype, rtol):
-    scale = 0.5 / math.sqrt(K)
-    a_np = np.random.uniform(0, scale, (M, K)).astype(np.float32)
-    b_np = np.random.uniform(0, scale, (K, N)).astype(np.float32)
-
-    a_mx = mx.array(a_np).astype(getattr(mx, dtype))
-    b_mx = mx.array(b_np).astype(getattr(mx, dtype))
-
-    a_pt = torch.from_numpy(a_np).to(dtype=getattr(torch, dtype), device="mps")
-    b_pt = torch.from_numpy(b_np).to(dtype=getattr(torch, dtype), device="mps")
-    torch.mps.synchronize()
-
-    torch_mean, torch_std = bench_torch(a_pt, b_pt)
-    mlx_mean, mlx_std = bench_mlx(a_mx, b_mx)
-
-    out_mx = (a_mx @ b_mx).astype(mx.float32)
-    out_pt = (a_pt @ b_pt).to(torch.float32).to("cpu").numpy(force=True)
-    check_correctness(out_mx, out_pt, rtol, M, N, K)
-
-    return mlx_mean, mlx_std, torch_mean, torch_std
-
-
-if __name__ == "__main__":
-    dtypes = ("bfloat16", "float16", "float32")
-
-    rtols = {
-        "float32": 1e-3,
-        "float16": 5e-3,
-        "bfloat16": 1e-2,
-    }
-
-    shapes = (
-        (2048, 2048, 10240),
-        (2048, 3072, 10240),
-        (3072, 3072, 10240),
-        (3072, 3072, 12288),
-        (3072, 4096, 12288),
-        (4096, 4096, 12288),
-        (4096, 4096, 18432),
-        (4096, 4096, 21504),
-        (4096, 6144, 21504),
-        (6144, 6144, 21504),
-    )
-
-    for dtype in dtypes:
-        print(f"\nPerformance ({dtype}):")
-        print(
-            f"{'M':>5s} {'N':>5s} {'K':>6s}  "
-            f"{'MLX (ms)':>15s}  {'Torch (ms)':>15s}  {'Speedup':>10s}"
-        )
-        print("-" * 80)
-
-        for M, N, K in shapes:
-            mlx_mean, mlx_std, torch_mean, torch_std = bench_gemm(
-                M, N, K, dtype, rtols[dtype]
-            )
-            speedup = torch_mean / mlx_mean
-
-            print(
-                f"{M:5d} {N:5d} {K:6d}  "
-                f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f}  "
-                f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f}  "
-                f"{speedup:8.2f}x"
-            )
@@ -1,7 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.

-from functools import partial
-
 import mlx.core as mx
 import mlx.nn as nn
 from time_utils import time_fn
@@ -12,71 +10,32 @@ def layer_norm(x, w, b, eps):
    x = x.astype(mx.float32)
    mu = mx.mean(x, -1, keepdims=True)
    v = mx.var(x, -1, keepdims=True)
-    y = (x - mu) * mx.rsqrt(v + eps)
-    if w is not None:
-        y = y * w
-    if b is not None:
-        y = y + b
-    return y
+    return (x - mu) * mx.rsqrt(v + eps) * w + b


-def time_layer_norm(N, dt):
-    L = 1024
+def time_layer_norm():
    f1 = lambda x, w, b, y: (layer_norm(x, w, b, 1e-5) * y).sum()
    f2 = lambda x, w, b, y: (mx.fast.layer_norm(x, w, b, 1e-5) * y).sum()
    g1 = mx.grad(f1, argnums=(0, 1, 2))
    g2 = mx.grad(f2, argnums=(0, 1, 2))

-    x = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    w = mx.random.uniform(shape=(N,)).astype(dt)
-    b = mx.random.uniform(shape=(N,)).astype(dt)
-    y = mx.random.uniform(shape=(8, L, N)).astype(dt)
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    b = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
    mx.eval(x, w, b, y)

-    def layer_norm_loop(f, x, w, b):
-        for _ in range(32):
-            x = f(x, w, b)
-        return x
-
-    time_fn(layer_norm_loop, partial(layer_norm, eps=1e-5), x, w, b)
-    time_fn(layer_norm_loop, partial(mx.fast.layer_norm, eps=1e-5), x, w, b)
-
-    def layer_norm_grad_loop(g, x, w, b):
+    def layer_norm_loop(g, x, w, b):
        gx, gw, gb = x, w, b
        for _ in range(32):
            gx, gw, gb = g(gx, gw, gb, y)
        return gx, gw, gb

-    time_fn(layer_norm_grad_loop, g1, x, w, b)
-    time_fn(layer_norm_grad_loop, g2, x, w, b)
-    time_fn(layer_norm_grad_loop, mx.compile(g1), x, w, b)
-    time_fn(layer_norm_grad_loop, mx.compile(g2), x, w, b)
-
-    f1 = lambda x, y: (layer_norm(x, None, None, 1e-5) * y).sum()
-    f2 = lambda x, y: (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0,))
-    g2 = mx.grad(f2, argnums=(0,))
-
-    x = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    w = mx.random.uniform(shape=(N,)).astype(dt)
-    b = mx.random.uniform(shape=(N,)).astype(dt)
-    y = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    mx.eval(x, w, b, y)
-
-    def layer_norm_grad_x_loop(g, x):
-        gx = x
-        for _ in range(32):
-            gx = g(gx, y)
-        return gx
-
-    time_fn(layer_norm_grad_x_loop, g1, x)
-    time_fn(layer_norm_grad_x_loop, g2, x)
-    time_fn(layer_norm_grad_x_loop, mx.compile(g1), x)
-    time_fn(layer_norm_grad_x_loop, mx.compile(g2), x)
+    time_fn(layer_norm_loop, g1, x, w, b)
+    time_fn(layer_norm_loop, g2, x, w, b)
+    time_fn(layer_norm_loop, mx.compile(g1), x, w, b)
+    time_fn(layer_norm_loop, mx.compile(g2), x, w, b)


 if __name__ == "__main__":
-    for dt in [mx.float32, mx.float16, mx.bfloat16]:
-        for n in [1024, 2048, 4096, 8192, 8192 + 1024]:
-            print(dt, n)
-            time_layer_norm(n, dt)
+    time_layer_norm()
@@ -1,236 +0,0 @@
-import math
-import os
-import platform
-import subprocess
-import time
-from copy import copy
-from functools import partial
-
-import matplotlib.pyplot as plt
-import mlx.core as mx
-import numpy as np
-import torch
-from matplotlib.ticker import FuncFormatter
-
-RESULTS_DIR = "./results"
-
-
-if not os.path.isdir(RESULTS_DIR):
-    os.mkdir(RESULTS_DIR)
-
-TORCH_DEVICE = torch.device(
-    "mps"
-    if torch.backends.mps.is_available()
-    else ("cuda" if torch.cuda.is_available() else "cpu")
-)
-
-
-def get_device_name():
-    if TORCH_DEVICE.type == "cuda":
-        try:
-            out = subprocess.check_output(
-                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
-                stderr=subprocess.DEVNULL,
-            )
-            return out.decode("utf-8").splitlines()[0].strip()
-        except Exception:
-            return "CUDA_GPU"
-    if TORCH_DEVICE.type == "mps":
-        try:
-            out = subprocess.check_output(
-                ["sysctl", "-n", "machdep.cpu.brand_string"],
-                stderr=subprocess.DEVNULL,
-            )
-            return out.decode("utf-8").strip()
-        except Exception:
-            return "Apple_Silicon"
-    return platform.processor() or platform.machine() or "CPU"
-
-
-DEVICE_NAME = get_device_name()
-
-
-N_WARMUP = 5
-N_ITER_BENCH = 50
-N_ITER_FUNC = 20
-
-VECTOR_LENGTHS = [4096 * (2**i) for i in range(12)]
-MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
-D_TYPES = ("float32", "float16")
-
-
-def _power_of_two_formatter(value, _position):
-    if value <= 0:
-        return ""
-    exponent = int(round(math.log2(value)))
-    if abs(value - (1 << exponent)) / value > 1e-6:
-        return f"{value:g}"
-    return f"$2^{{{exponent}}}$"
-
-
-def torch_sync():
-    if TORCH_DEVICE.type == "cuda":
-        torch.cuda.synchronize()
-    elif TORCH_DEVICE.type == "mps":
-        torch.mps.synchronize()
-
-
-def masked_scatter_mlx(self_arr, mask_arr, src_arr):
-    outs = []
-    for _ in range(N_ITER_FUNC):
-        out = copy(self_arr)
-        out[mask_arr] = src_arr
-        outs.append(out)
-    mx.eval(outs)
-    return outs
-
-
-@torch.no_grad()
-def masked_scatter_torch(self_tensor, mask_tensor, src_tensor):
-    outs = []
-    for _ in range(N_ITER_FUNC):
-        out = self_tensor.clone()
-        out.masked_scatter_(mask_tensor, src_tensor)
-        outs.append(out)
-    torch_sync()
-    return outs
-
-
-def measure(fn):
-    for _ in range(N_WARMUP):
-        fn()
-    start = time.perf_counter_ns()
-    for _ in range(N_ITER_BENCH):
-        fn()
-    end = time.perf_counter_ns()
-    return (end - start) * 1e-9
-
-
-def bytes_touched(length, true_count, item_size):
-    mask_bytes = length
-    self_bytes = length * item_size * 2  # read + write
-    src_bytes = true_count * item_size
-    return (mask_bytes + self_bytes + src_bytes) * N_ITER_FUNC * N_ITER_BENCH
-
-
-def build_case(length, density, np_dtype, torch_dtype):
-    true_count = max(1, int(round(length * density)))
-
-    rng = np.random.default_rng()
-    self_np = rng.normal(0.0, 1.0, length).astype(np_dtype)
-    mask_np = np.zeros(length, dtype=bool)
-    mask_np[:true_count] = True
-    rng.shuffle(mask_np)
-    src_np = rng.normal(0.0, 1.0, true_count).astype(np_dtype)
-
-    self_mlx = mx.array(self_np)
-    mask_mlx = mx.array(mask_np)
-    src_mlx = mx.array(src_np)
-
-    self_torch = torch.from_numpy(self_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
-    mask_torch = torch.from_numpy(mask_np).to(device=TORCH_DEVICE)
-    src_torch = torch.from_numpy(src_np).to(device=TORCH_DEVICE, dtype=torch_dtype)
-
-    # Correctness check once per configuration
-    mx_out = mx.array(self_np)
-    mx_out[mask_mlx] = src_mlx
-    mx.eval(mx_out)
-    torch_out = self_torch.clone()
-    torch_out.masked_scatter_(mask_torch, src_torch)
-
-    atol = 5e-3 if np_dtype == np.float16 else 1e-5
-    if not np.allclose(np.array(mx_out), torch_out.cpu().numpy(), atol=atol):
-        raise AssertionError("masked_scatter results diverged between MLX and Torch")
-
-    return (self_mlx, mask_mlx, src_mlx, self_torch, mask_torch, src_torch, true_count)
-
-
-def bench_case(length, density, dtype):
-    np_dtype = getattr(np, dtype)
-    torch_dtype = getattr(torch, dtype)
-    (
-        self_mlx,
-        mask_mlx,
-        src_mlx,
-        self_torch,
-        mask_torch,
-        src_torch,
-        true_count,
-    ) = build_case(length, density, np_dtype, torch_dtype)
-
-    time_mlx = measure(partial(masked_scatter_mlx, self_mlx, mask_mlx, src_mlx))
-    time_torch = measure(
-        partial(masked_scatter_torch, self_torch, mask_torch, src_torch)
-    )
-
-    total_bytes = bytes_touched(length, true_count, np_dtype().itemsize)
-    bytes_per_gb = float(1024**3)
-    mlx_gbps = (total_bytes / bytes_per_gb) / time_mlx
-    torch_gbps = (total_bytes / bytes_per_gb) / time_torch
-
-    return time_mlx, time_torch, mlx_gbps, torch_gbps
-
-
-def plot_density(ax_perf, ax_speedup, density, dtype):
-    mlx_gbps = []
-    torch_gbps = []
-    mlx_times = []
-    torch_times = []
-
-    for length in VECTOR_LENGTHS:
-        t_mlx, t_torch, gbps_mlx, gbps_torch = bench_case(length, density, dtype)
-        mlx_gbps.append(gbps_mlx)
-        torch_gbps.append(gbps_torch)
-        mlx_times.append(t_mlx)
-        torch_times.append(t_torch)
-
-    ax_perf.plot(VECTOR_LENGTHS, mlx_gbps, "tab:blue", label="MLX")
-    ax_perf.plot(VECTOR_LENGTHS, torch_gbps, "tab:red", label="Torch")
-    ax_perf.set_xscale("log", base=2)
-    ax_perf.set_xticks(VECTOR_LENGTHS)
-    formatter = FuncFormatter(_power_of_two_formatter)
-    ax_perf.xaxis.set_major_formatter(formatter)
-    ax_perf.set_title(f"density={density:.2f}")
-    ax_perf.set_ylabel("GB/s")
-    ax_perf.grid(True, which="both", linestyle=":", alpha=0.4)
-    ax_perf.legend()
-
-    speedup = np.array(torch_times) / np.array(mlx_times)
-    ax_speedup.plot(VECTOR_LENGTHS, speedup, "tab:green")
-    ax_speedup.axhline(1.0, color="tab:gray", linestyle="--")
-    ax_speedup.set_xscale("log", base=2)
-    ax_speedup.set_xticks(VECTOR_LENGTHS)
-    ax_speedup.xaxis.set_major_formatter(formatter)
-    ax_speedup.set_ylabel("Speedup (Torch_t / MLX_t)")
-    ax_speedup.grid(True, which="both", linestyle=":", alpha=0.4)
-
-
-def main():
-    for dtype in D_TYPES:
-        fig, axs = plt.subplots(
-            len(MASK_DENSITIES),
-            2,
-            figsize=(10, 12),
-            layout="constrained",
-            sharex=True,
-        )
-
-        for i, density in enumerate(MASK_DENSITIES):
-            plot_density(axs[i][0], axs[i][1], density, dtype)
-            axs[i][0].set_xlabel("vector length")
-            axs[i][1].set_xlabel("vector length")
-
-        fig.suptitle(
-            f"{DEVICE_NAME.replace('Apple ', '')} ({TORCH_DEVICE.type}) | dtype={dtype}"
-        )
-        output_path = os.path.join(
-            RESULTS_DIR,
-            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.png",
-        )
-        fig.savefig(output_path)
-        print(f"Saved benchmark image: {output_path}")
-        plt.close(fig)
-
-
-if __name__ == "__main__":
-    main()
@@ -9,10 +9,7 @@ def rms_norm(x, w, eps):
    ot = x.dtype
    x = x.astype(mx.float32)
    n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
-    y = (x * n).astype(ot)
-    if w is not None:
-        y = y * w
-    return y
+    return (x * n).astype(ot) * w


 def time_rms_norm():
@@ -37,27 +34,6 @@ def time_rms_norm():
    time_fn(rms_norm_loop, mx.compile(g1), x, w)
    time_fn(rms_norm_loop, mx.compile(g2), x, w)

-    f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum()
-    f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0,))
-    g2 = mx.grad(f2, argnums=(0,))
-
-    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
-    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    mx.eval(x, w, y)
-
-    def rms_norm_loop(g, x):
-        gx = x
-        for _ in range(32):
-            gx = g(gx, y)
-        return gx
-
-    time_fn(rms_norm_loop, g1, x)
-    time_fn(rms_norm_loop, g2, x)
-    time_fn(rms_norm_loop, mx.compile(g1), x)
-    time_fn(rms_norm_loop, mx.compile(g2), x)
-

 if __name__ == "__main__":
    time_rms_norm()
@@ -28,34 +28,11 @@ def bench(f, *args):
    return (e - s) * 1e-9


-def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype):
-    np_dtype = getattr(np, dtype)
-
-    shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D)
-    shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D)
-
-    scale = 1.0 / math.sqrt(D)
-
-    q_np = np.random.normal(0.0, 1.0, shape_q).astype(np_dtype)
-    k_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-    v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-
-    q_mx = mx.array(q_np)
-    k_mx = mx.array(k_np)
-    v_mx = mx.array(v_np)
-
-    if mask is not None:
-        if mask == "additive":
-            mask_np = np.random.normal(0.0, 1.0, (B, qH, qL, kL)).astype(np_dtype)
-            mask = mx.array(mask_np)
-        elif mask == "bool":
-            mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
-            mask = mx.array(mask_np)
-
-    return q_mx, k_mx, v_mx, scale, mask
+def mlx_sdpa_fused_inner(q, k, v, scale):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None)


-def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
+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]
@@ -64,7 +41,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):

    B = q.shape[0]
    L = q.shape[2]
-    kL = k.shape[2]

    if n_repeats > 1:
        q = mx.reshape(q, [B, n_kv_heads, n_repeats, L, -1])
@@ -72,27 +48,10 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
        v = mx.expand_dims(v, 2)

    scores = q @ mx.swapaxes(k, -1, -2)
-
-    if mask is not None:
-
-        if mask == "causal":
-            q_offset = max(0, kL - L)
-            q_indices = mx.arange(q_offset, q_offset + L)
-            k_indices = mx.arange(kL)
-            mask = q_indices[:, None] >= k_indices[None]
-
-        if n_repeats > 1 and mask.ndim >= 3:
-            if mask.shape[-3] == 1:
-                mask = mx.expand_dims(mask, -3)
-            else:
-                mask = mx.unflatten(mask, -3, (n_kv_heads, n_repeats))
-
-        if mask.dtype == mx.bool_:
-            scores = mx.where(mask, scores, -np.float32(np.inf))
-        else:
-            scores += mask
-
-    scores = mx.softmax(scores, axis=-1, precise=True)
+    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:
@@ -101,55 +60,74 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
    return out


-def mlx_fused_attn(q, k, v, scale, mask):
-    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention(f, q, k, v, scale, mask=None, transpose=False):
-    if transpose:
-        q_t = mx.transpose(q, (0, 2, 1, 3))
-        k_t = mx.transpose(k, (0, 2, 1, 3))
-        v_t = mx.transpose(v, (0, 2, 1, 3))
-        o_t = f(q_t, k_t, v_t, scale=scale, mask=mask)
-        return mx.transpose(o_t, (0, 2, 1, 3))
-    else:
-        return f(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention_bench(f, q, k, v, scale, mask=None, transpose=False):
+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):
-        q_out = do_attention(f, q_out, k, v, scale, mask=mask, transpose=transpose)
+        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 bench_shape(
-    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None
-):
-    q_mx, k_mx, v_mx, scale, mask = prepare_inputs(
-        B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, mask_in, transpose, dtype
+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)
    )

-    time_mlx_unfused = bench(
-        do_attention_bench, mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
-    time_mlx_fused = bench(
-        do_attention_bench, mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    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)

-    o_mlx_fused = do_attention(mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose)
-    o_mlx_unfused = do_attention(
-        mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    scale = math.sqrt(1.0 / head_dim)

-    atol = 1e-5 if dtype == "float32" else 2e-4
+    q_mx = mx.array(q_np)
+    k_mx = mx.array(k_np)
+    v_mx = mx.array(v_np)

-    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol, rtol=atol):
+    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}, mask: {mask_in}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
+            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
@@ -173,57 +151,39 @@ if __name__ == "__main__":
          (  1,   128,   128,       64,   32,    32),
          (  1,   256,   256,       64,   32,    32),
          (  1,   512,   512,       64,   32,    32),
-          (  1,  1024,  1024,       64,   32,     8),
-          (  1,  2048,  2048,       64,   32,     8),
-          (  1,  4096,  4096,       64,   32,     8),
-          (  1,  4096,  5000,       64,   32,     8),
-          (  1,  2048,  32121,      64,   32,     8),
+          (  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,     8),
-          (  1,  2048,  2048,       80,   32,     8),
-          (  1,  4096,  4096,       80,   32,     8),
-          (  1,  4096,  5000,       80,   32,     8),
-          (  1,  2048,  32121,      80,   32,     8),
+          (  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,     8),
-          (  1,  2048,  2048,      128,   32,     8),
-          (  1,  4096,  4096,      128,   32,     8),
-          (  1,  4096,  5000,      128,   32,     8),
-          (  1,  2048,  32121,     128,   32,     8),
+          (  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

-    masks = [None, "bool", "causal"]
-
-    print(
-        "  B,   qsl,   ksl, hdim, n_qh, n_kvh, t,   dtype,     mask, t_unfs, t_fuse, diff%"
-    )
+    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:
-                for mask_in in masks:
-                    time_mlx_fused, time_mlx_unfused = bench_shape(
-                        B,
-                        qsl,
-                        ksl,
-                        head_dim,
-                        n_q_heads,
-                        n_kv_heads,
-                        dtype,
-                        transpose,
-                        mask_in,
-                    )
-                    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:1d}, {dtype}, {str(mask_in):>8}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
-                    )
+                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}%"
+                )
@@ -1,95 +1,94 @@
-import argparse
-import math
-
 import mlx.core as mx
+import numpy as np
+from mlx.utils import tree_map
 from time_utils import time_fn

-L = 16384
+L = 32768
 H = 32
 H_k = H // 4
 D = 128
-V = 128
 dtype = mx.float16
-loops = 10
+bits = 8
+
+loops = 20


-def upproject(x, w):
-    if w is None:
-        return x
-    else:
-        return x @ w.T
-
-
-def attention(q, k, v, mask=None, w=None):
-    def _sdpa(q, k, v):
+def attention(q, k, v):
+    for _ in range(loops):
        B, Hq, L, D = q.shape
        _, Hk, S, _ = k.shape
-        _, _, _, V = v.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)
-        if mask is not None:
-            m = mx.broadcast_to(mask, (B, Hq, L, S)).reshape(B, Hk, Hq // Hk, L, S)
-            s = mx.where(m, s, mx.finfo(s.dtype).min)
+        ke = k[:, :, None, :, :]
+        ve = v[:, :, None, :, :]
+        s = q @ ke.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, V)
-
-    for i in range(loops):
-        q = _sdpa(q, k, v)
-        q = upproject(q, w)
+        q = p @ ve
+        q = q.reshape(B, Hq, L, D)
    return q


-def sdpa(q, k, v, mask=None, w=None):
-    for i in range(loops):
-        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
-        q = upproject(q, w)
+def sdpa(q, k, v):
+    for _ in range(loops):
+        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=None)
    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, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mx.eval(q, k, v, w)
-    time_fn(attention, q, k, v, w=w)
+def quant_sdpa(q, k, v, bits=4):
+    for _ in range(loops):
+        q = mx.fast.quantized_scaled_dot_product_attention(
+            q, *k, *v, scale=1.0, mask=None, bits=bits
+        )
+    return q


-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, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mx.eval(q, k, v, w)
-    time_fn(sdpa, q, k, v, w=w)
+def quant_attention(q, k, v, bits=4):
+    for _ in range(loops):
+        B, Hq, L, D = q.shape
+        Hk = k[0].shape[1]
+
+        q = q.reshape((B, Hk, Hq // Hk, L, D))
+        ke = tree_map(lambda x: mx.expand_dims(x, axis=2), k)
+        ve = tree_map(lambda x: mx.expand_dims(x, axis=2), v)
+
+        scores = mx.quantized_matmul(q, *ke, transpose=True, bits=bits)
+        scores = mx.softmax(scores, axis=-1)
+
+        q = mx.quantized_matmul(scores, *ve, transpose=False, bits=bits)
+        q = q.reshape((B, Hq, L, D))
+    return q


-def time_self_attention_sdpa_with_mask():
-    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, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mask = mx.full((L,), True)
-    mask[L // 2 :] = False
-    mx.eval(q, k, v, mask, w)
+def time_self_attention_primitives(q, k, v):
+    time_fn(attention, q, k, v)

-    def sdpa_mask(*args):
-        return sdpa(*args, mask=mask, w=w)

-    def attention_mask(*args):
-        return attention(*args, mask=mask, w=w)
+def time_self_attention_sdpa(q, k, v):
+    time_fn(sdpa, q, k, v)

-    time_fn(attention_mask, q, k, v)
-    time_fn(sdpa_mask, q, k, v)
+
+def time_self_attention_quant_sdpa(q, k, v, bits=4):
+    time_fn(quant_sdpa, q, k, v, bits)
+
+
+def time_self_attention_quant_primitives(q, k, v, bits=4):
+    time_fn(quant_attention, q, k, v, bits)


 if __name__ == "__main__":
-    time_self_attention_sdpa()
-    time_self_attention_primitives()
-    time_self_attention_sdpa_with_mask()
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D), dtype=dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D), dtype=dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, D), dtype=dtype)
+    mx.eval(q, k, v)
+
+    k_quant = mx.quantize(k, bits=bits)
+    v_quant = mx.quantize(v, bits=bits)
+    mx.eval(k_quant, v_quant)
+
+    k = mx.dequantize(*k_quant, bits=bits)
+    v = mx.dequantize(*v_quant, bits=bits)
+
+    time_self_attention_sdpa(q, k, v)
+    time_self_attention_quant_sdpa(q, k_quant, v_quant, bits)
+    time_self_attention_primitives(q, k, v)
+    time_self_attention_quant_primitives(q, k_quant, v_quant, bits)
@@ -1,209 +0,0 @@
-# Copyright © 2026 Apple Inc.
-
-import argparse
-import time
-
-import mlx.core as mx
-import numpy as np
-
-MLX_DTYPES = {
-    "float16": mx.float16,
-    "bfloat16": mx.bfloat16,
-    "float32": mx.float32,
-}
-
-
-def parse_cases(cases):
-    parsed = []
-    for spec in cases.split(","):
-        m, n, k, s = [int(x) for x in spec.split("x")]
-        parsed.append((m, n, k, s))
-    return parsed
-
-
-def make_segments(k, num_segments, pattern, seed):
-    if pattern == "equal":
-        cuts = np.linspace(0, k, num_segments + 1, dtype=np.int64)
-    else:
-        rng = np.random.default_rng(seed)
-        cuts = rng.integers(0, k + 1, size=(num_segments - 1,), dtype=np.int64)
-        cuts = np.sort(cuts)
-        cuts = np.concatenate(([0], cuts, [k]))
-    return np.stack([cuts[:-1], cuts[1:]], axis=1).astype(np.uint32)
-
-
-def numpy_segmented_mm_ref(a, b, segments):
-    """Ground-truth reference in float64."""
-    out = []
-    for start, end in segments:
-        out.append(a[:, start:end] @ b[start:end, :])
-    return np.stack(out, axis=0)
-
-
-def mlx_segmented_mm_loop(a, b, segments):
-    """MLX loop-of-matmuls baseline."""
-    segments_list = segments.tolist()
-    out = []
-    for start, end in segments_list:
-        out.append(a[:, start:end] @ b[start:end, :])
-    return mx.stack(out, axis=0)
-
-
-def bench_mlx(a, b, segments, warmup, iters):
-    for _ in range(warmup):
-        y = mx.segmented_mm(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = mx.segmented_mm(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-    end = time.perf_counter()
-    return (end - start) * 1e3 / iters
-
-
-def bench_mlx_loop(a, b, segments, warmup, iters):
-    for _ in range(warmup):
-        y = mlx_segmented_mm_loop(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-
-    start = time.perf_counter()
-    for _ in range(iters):
-        y = mlx_segmented_mm_loop(a, b, segments)
-        mx.eval(y)
-    mx.synchronize()
-    end = time.perf_counter()
-    return (end - start) * 1e3 / iters
-
-
-def print_table(headers, rows):
-    widths = [len(h) for h in headers]
-    for row in rows:
-        for i, cell in enumerate(row):
-            widths[i] = max(widths[i], len(cell))
-
-    def fmt_row(row):
-        return (
-            "| "
-            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
-            + " |"
-        )
-
-    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
-    print(fmt_row(headers))
-    print(sep)
-    for row in rows:
-        print(fmt_row(row))
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--cases",
-        default=(
-            "128x128x1024x16,"
-            "128x128x1024x32,"
-            "256x256x2048x16,"
-            "512x512x4096x32,"
-            "1024x1024x4096x32,"
-            "1024x1024x8192x64"
-        ),
-        help="Comma-separated MxNxKxS list.",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="float32",
-        choices=["float16", "bfloat16", "float32"],
-    )
-    parser.add_argument("--warmup", type=int, default=10)
-    parser.add_argument("--iters", type=int, default=50)
-    parser.add_argument(
-        "--segments",
-        choices=["equal", "random"],
-        default="random",
-        help="Segment generation pattern.",
-    )
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--no-check", action="store_true")
-    args = parser.parse_args()
-
-    mlx_dtype = MLX_DTYPES[args.dtype]
-
-    print(
-        f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}"
-    )
-
-    headers = [
-        "Case",
-        "MLX ms",
-        "Loop ms",
-        "Speedup",
-        "MLX err",
-        "Loop err",
-    ]
-    rows = []
-
-    cases = parse_cases(args.cases)
-    for idx, (m, n, k, s) in enumerate(cases):
-        rng = np.random.default_rng(args.seed + idx)
-        a_np = rng.standard_normal((m, k)).astype(np.float32)
-        b_np = rng.standard_normal((k, n)).astype(np.float32)
-        seg_np = make_segments(k, s, args.segments, args.seed + idx)
-
-        a_mx = mx.array(a_np, dtype=mlx_dtype)
-        b_mx = mx.array(b_np, dtype=mlx_dtype)
-        seg_mx = mx.array(seg_np, dtype=mx.uint32)
-        mx.eval(a_mx, b_mx, seg_mx)
-
-        mlx_err_str = ""
-        loop_err_str = ""
-        if not args.no_check:
-            y_mlx = mx.segmented_mm(a_mx, b_mx, seg_mx)
-            y_loop = mlx_segmented_mm_loop(a_mx, b_mx, seg_mx)
-            mx.eval(y_mlx, y_loop)
-
-            if args.dtype == "float32":
-                ref = numpy_segmented_mm_ref(
-                    a_np.astype(np.float64),
-                    b_np.astype(np.float64),
-                    seg_np.tolist(),
-                )
-                mlx_err = np.max(np.abs(np.array(y_mlx, dtype=np.float64) - ref))
-                loop_err = np.max(np.abs(np.array(y_loop, dtype=np.float64) - ref))
-            else:
-                a_mx_f32 = mx.array(a_np, dtype=mx.float32)
-                b_mx_f32 = mx.array(b_np, dtype=mx.float32)
-                ref = mx.segmented_mm(a_mx_f32, b_mx_f32, seg_mx)
-                mx.eval(ref)
-                mlx_err = float(mx.max(mx.abs(ref - y_mlx.astype(mx.float32))).item())
-                loop_err = float(mx.max(mx.abs(ref - y_loop.astype(mx.float32))).item())
-            mlx_err_str = f"{mlx_err:.2e}"
-            loop_err_str = f"{loop_err:.2e}"
-
-        t_mlx = bench_mlx(a_mx, b_mx, seg_mx, args.warmup, args.iters)
-        t_loop = bench_mlx_loop(a_mx, b_mx, seg_mx, args.warmup, args.iters)
-        ratio = t_loop / t_mlx if t_mlx > 0 else float("inf")
-        rows.append(
-            [
-                f"{m}x{n}x{k}x{s}",
-                f"{t_mlx:.3f}",
-                f"{t_loop:.3f}",
-                f"{ratio:.2f}x",
-                mlx_err_str,
-                loop_err_str,
-            ]
-        )
-
-    print_table(headers, rows)
-    if not args.no_check:
-        if args.dtype == "float32":
-            print("err: max|result - numpy_fp64_ref|")
-        else:
-            print("err: max|result - own_fp32_result|")
-
-
-if __name__ == "__main__":
-    main()
@@ -51,20 +51,6 @@ def time_maximum():
    time_fn(mx.maximum, a, b)


-def time_max():
-    a = mx.random.uniform(shape=(32, 1024, 1024))
-    a[1, 1] = mx.nan
-    mx.eval(a)
-    time_fn(mx.max, a, 0)
-
-
-def time_min():
-    a = mx.random.uniform(shape=(32, 1024, 1024))
-    a[1, 1] = mx.nan
-    mx.eval(a)
-    time_fn(mx.min, a, 0)
-
-
 def time_negative():
    a = mx.random.uniform(shape=(10000, 1000))
    mx.eval(a)
@@ -122,8 +108,6 @@ if __name__ == "__main__":

    time_add()
    time_matmul()
-    time_min()
-    time_max()
    time_maximum()
    time_exp()
    time_negative()
@@ -1,109 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import argparse
-
-import mlx.core as mx
-import torch
-from time_utils import measure_runtime
-
-
-def benchmark_slice_update_mlx(dst_shape, slice_shape, slice_range, dtype, iters=10):
-    def slice_update(arguments):
-        for i in range(iters):
-            arguments["dst"] = (
-                arguments["dst"].at[slice_range].add(arguments["updates"])
-            )
-        mx.eval(arguments)
-
-    dtype = getattr(mx, dtype)
-    arguments = {
-        "dst": mx.random.normal(dst_shape).astype(dtype),
-        "updates": mx.random.normal(slice_shape).astype(dtype),
-    }
-
-    runtime = measure_runtime(slice_update, arguments=arguments)
-    bytes_processed = (
-        arguments["dst"][slice_range].nbytes * 2 + arguments["updates"].nbytes
-    ) * iters
-    bandwidth_gb_s = bytes_processed / runtime / 1e6
-    return runtime, bandwidth_gb_s
-
-
-def benchmark_slice_update_torch(
-    dst_shape, slice_shape, slice_range, device, dtype, iters=10
-):
-    def slice_update(dst, updates, slice_range):
-        for i in range(iters):
-            dst[slice_range] = dst[slice_range] + updates
-        if device == torch.device("mps"):
-            torch.mps.synchronize()
-
-    dtype = getattr(torch, dtype)
-    updates = torch.randn(slice_shape, dtype=dtype).to(device)
-    dst = torch.randn(dst_shape, dtype=dtype).to(device)
-
-    runtime = measure_runtime(
-        slice_update, dst=dst, updates=updates, slice_range=slice_range
-    )
-    bytes_processed = (dst[slice_range].nbytes * 2 + updates.nbytes) * iters
-    bandwidth_gb_s = bytes_processed / runtime / 1e6
-    return runtime, bandwidth_gb_s
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser("Slice update benchmarks.")
-    parser.add_argument("--cpu", action="store_true", help="Use the CPU.")
-    args = parser.parse_args()
-
-    if args.cpu:
-        mx.set_default_device(mx.cpu)
-        device = torch.device("cpu")
-    elif torch.mps.is_available():
-        device = torch.device("mps")
-    elif torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        raise ValueError()
-
-    dtypes = ["float32", "bfloat16"]
-
-    test_cases = [
-        ((10_000_000,), slice(0, 1_000_000), (1_000_000,)),
-        ((100_000,), slice(10_000, 20_000), (10_000,)),
-        ((1000, 64), slice(100, 200), (100, 64)),
-        ((100, 100, 64), slice(20, 40), (20, 100, 64)),
-        (
-            (2048, 2048, 128),
-            (slice(500, 1500), slice(200, 1200), slice(32, 96)),
-            (1000, 1000, 64),
-        ),
-        (
-            (2048, 2048, 128),
-            (slice(1800, 1850), slice(100, 200), slice(64, 128)),
-            (50, 100, 64),
-        ),
-        (
-            (2048, 2048, 128),
-            (slice(1000, 1010), slice(1000, 1010), slice(64, 128)),
-            (10, 10, 64),
-        ),
-    ]
-
-    print(
-        f"{'Dtype':<12} {'Dst Shape':<25} {'Update Shape':<20} "
-        f"{'MLX (ms)':<12} {'MLX GB/s':<12} {'Torch (ms)':<12} {'Torch GB/s':<12}"
-    )
-    print("-" * 110)
-
-    for dtype in dtypes:
-        for dst_shape, slice_range, update_shape in test_cases:
-            mlx_time, mlx_bw = benchmark_slice_update_mlx(
-                dst_shape, update_shape, slice_range, dtype
-            )
-            torch_time, torch_bw = benchmark_slice_update_torch(
-                dst_shape, update_shape, slice_range, device, dtype
-            )
-            print(
-                f"{dtype:<12} {str(dst_shape):<25} {str(update_shape):<20} "
-                f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} {torch_bw:<12.2f}"
-            )
@@ -1,55 +0,0 @@
-import time
-
-import mlx.core as mx
-
-rank = mx.distributed.init().rank()
-
-
-def timeit(fn, a):
-
-    # warmup
-    for _ in range(5):
-        mx.eval(fn(a))
-
-    its = 10
-    tic = time.perf_counter()
-    for _ in range(its):
-        mx.eval(fn(a))
-    toc = time.perf_counter()
-    ms = 1000 * (toc - tic) / its
-    return ms
-
-
-def all_reduce_benchmark():
-    a = mx.ones((5, 5), mx.int32)
-
-    its_per_eval = 100
-
-    def fn(x):
-        for _ in range(its_per_eval):
-            x = mx.distributed.all_sum(x)
-            x = x - 1
-        return x
-
-    ms = timeit(fn, a) / its_per_eval
-    if rank == 0:
-        print(f"All Reduce: time per iteration {ms:.6f} (ms)")
-
-
-def all_gather_benchmark():
-    a = mx.ones((5, 5), mx.int32)
-    its_per_eval = 100
-
-    def fn(x):
-        for _ in range(its_per_eval):
-            x = mx.distributed.all_gather(x)[0]
-        return x
-
-    ms = timeit(fn, a) / its_per_eval
-    if rank == 0:
-        print(f"All gather: time per iteration {ms:.6f} (ms)")
-
-
-if __name__ == "__main__":
-    all_reduce_benchmark()
-    all_gather_benchmark()
@@ -31,8 +31,8 @@ def measure_runtime(fn, **kwargs):
    for _ in range(5):
        fn(**kwargs)

-    tic = time.perf_counter()
+    tic = time.time()
    iters = 100
    for _ in range(iters):
        fn(**kwargs)
-    return (time.perf_counter() - tic) * 1000 / iters
+    return (time.time() - tic) * 1000 / iters
@@ -1,177 +0,0 @@
-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# Modified from
-# https://github.com/NVIDIA/cudnn-frontend/blob/main/cmake/cuDNN.cmake
-
-# Return the last file matching the pattern.
-function(find_file_glob VAR PATTERN)
-  file(GLOB _RESULT "${PATTERN}")
-  if(_RESULT)
-    list(LENGTH ${_RESULT} _RESULT_LENGTH)
-    if(_RESULT_LENGTH GREATER 0)
-      list(GET ${_RESULT} -1 _RESULT)
-    endif()
-    set(${VAR}
-        "${_RESULT}"
-        PARENT_SCOPE)
-  endif()
-endfunction()
-
-# Find the dir including the "cudnn.h" file.
-find_path(
-  CUDNN_INCLUDE_DIR cudnn.h
-  HINTS ${CUDNN_INCLUDE_PATH} ${CUDAToolkit_INCLUDE_DIRS}
-  PATH_SUFFIXES include OPTIONAL)
-
-# Glob searching "cudnn.h" for Windows.
-if(WIN32 AND NOT CUDNN_INCLUDE_DIR)
-  find_file_glob(
-    CUDNN_H_PATH
-    "C:/Program Files/NVIDIA/CUDNN/*/include/${CUDAToolkit_VERSION_MAJOR}.*/cudnn.h"
-  )
-  if(CUDNN_H_PATH)
-    get_filename_component(CUDNN_INCLUDE_DIR "${CUDNN_H_PATH}" DIRECTORY)
-  endif()
-endif()
-
-if(NOT CUDNN_INCLUDE_DIR)
-  message(
-    FATAL_ERROR
-      "Unable to find cudnn.h, please make sure cuDNN is installed and pass CUDNN_INCLUDE_PATH to cmake."
-  )
-endif()
-
-# Get cudnn version.
-file(READ "${CUDNN_INCLUDE_DIR}/cudnn_version.h" cudnn_version_header)
-string(REGEX MATCH "#define CUDNN_MAJOR [1-9]+" macrodef
-             "${cudnn_version_header}")
-string(REGEX MATCH "[1-9]+" CUDNN_MAJOR_VERSION "${macrodef}")
-
-# Function for searching library files.
-function(find_cudnn_library NAME)
-  if(NOT "${ARGV1}" STREQUAL "OPTIONAL")
-    set(_CUDNN_REQUIRED TRUE)
-  else()
-    set(_CUDNN_REQUIRED FALSE)
-  endif()
-
-  find_library(
-    ${NAME}_LIBRARY
-    NAMES ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}" NAMES_PER_DIR
-    HINTS ${CUDNN_LIBRARY_PATH} ${CUDAToolkit_LIBRARY_DIR}
-    PATH_SUFFIXES lib64 lib/x64 lib OPTIONAL)
-
-  if(WIN32 AND NOT ${NAME}_LIBRARY)
-    find_file_glob(
-      ${NAME}_LIBRARY
-      "C:/Program Files/NVIDIA/CUDNN/*/lib/${CUDAToolkit_VERSION_MAJOR}.*/x64/${NAME}.lib"
-    )
-  endif()
-
-  if(NOT ${NAME}_LIBRARY AND ${_CUDNN_REQUIRED})
-    message(
-      FATAL_ERROR
-        "Unable to find ${NAME}, please make sure cuDNN is installed and pass CUDNN_LIBRARY_PATH to cmake."
-    )
-  endif()
-
-  if(${NAME}_LIBRARY)
-    add_library(CUDNN::${NAME} UNKNOWN IMPORTED)
-    set_target_properties(
-      CUDNN::${NAME}
-      PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR}
-                 IMPORTED_LOCATION ${${NAME}_LIBRARY})
-    set(${NAME}_LIBRARY
-        "${${NAME}_LIBRARY}"
-        PARENT_SCOPE)
-  else()
-    message(STATUS "${NAME} not found.")
-  endif()
-endfunction()
-
-# Search for the main cudnn library.
-find_cudnn_library(cudnn)
-
-include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(CUDNN REQUIRED_VARS CUDNN_INCLUDE_DIR
-                                                      cudnn_LIBRARY)
-
-if(CUDNN_INCLUDE_DIR AND cudnn_LIBRARY)
-  set(CUDNN_FOUND
-      ON
-      CACHE INTERNAL "cuDNN Library Found")
-else()
-  set(CUDNN_FOUND
-      OFF
-      CACHE INTERNAL "cuDNN Library Not Found")
-endif()
-
-# Find out all the DLL files for Windows.
-if(WIN32 AND cudnn_LIBRARY)
-  get_filename_component(CUDNN_BIN_DIR "${cudnn_LIBRARY}" DIRECTORY)
-  string(REPLACE "/lib/" "/bin/" CUDNN_BIN_DIR "${CUDNN_BIN_DIR}")
-  file(
-    GLOB CUDNN_DLL_NAMES
-    RELATIVE "${CUDNN_BIN_DIR}"
-    "${CUDNN_BIN_DIR}/*.dll")
-endif()
-
-# Create an interface library that users can link with.
-add_library(CUDNN::cudnn_all INTERFACE IMPORTED)
-target_link_libraries(CUDNN::cudnn_all INTERFACE CUDNN::cudnn)
-target_include_directories(
-  CUDNN::cudnn_all INTERFACE $<INSTALL_INTERFACE:include>
-                             $<BUILD_INTERFACE:${CUDNN_INCLUDE_DIR}>)
-
-# Add other components of cudnn.
-if(CUDNN_MAJOR_VERSION EQUAL 8)
-  find_cudnn_library(cudnn_adv_infer)
-  find_cudnn_library(cudnn_adv_train)
-  find_cudnn_library(cudnn_cnn_infer)
-  find_cudnn_library(cudnn_cnn_train)
-  find_cudnn_library(cudnn_ops_infer)
-  find_cudnn_library(cudnn_ops_train)
-
-  target_link_libraries(
-    CUDNN::cudnn_all
-    INTERFACE CUDNN::cudnn_adv_train CUDNN::cudnn_ops_train
-              CUDNN::cudnn_cnn_train CUDNN::cudnn_adv_infer
-              CUDNN::cudnn_cnn_infer CUDNN::cudnn_ops_infer)
-
-elseif(CUDNN_MAJOR_VERSION EQUAL 9)
-  find_cudnn_library(cudnn_graph)
-  find_cudnn_library(cudnn_engines_runtime_compiled)
-  find_cudnn_library(cudnn_ops OPTIONAL)
-  find_cudnn_library(cudnn_cnn OPTIONAL)
-  find_cudnn_library(cudnn_adv OPTIONAL)
-  find_cudnn_library(cudnn_engines_precompiled OPTIONAL)
-  find_cudnn_library(cudnn_heuristic OPTIONAL)
-
-  target_link_libraries(
-    CUDNN::cudnn_all
-    INTERFACE CUDNN::cudnn_graph
-              CUDNN::cudnn_engines_runtime_compiled
-              CUDNN::cudnn_ops
-              CUDNN::cudnn_cnn
-              CUDNN::cudnn_adv
-              CUDNN::cudnn_engines_precompiled
-              CUDNN::cudnn_heuristic)
-endif()
@@ -1,54 +0,0 @@
-# FindNCCL.cmake This module finds the NVIDIA NCCL library and its include
-# directories.
-
-set(NCCL_ROOT_DIR
-    $ENV{NCCL_ROOT_DIR}
-    CACHE PATH "Folder contains NVIDIA NCCL")
-
-find_path(
-  NCCL_INCLUDE_DIRS
-  NAMES nccl.h
-  HINTS ${NCCL_INCLUDE_DIR} ${NCCL_ROOT_DIR} ${NCCL_ROOT_DIR}/include
-        ${CUDA_TOOLKIT_ROOT_DIR}/include)
-
-if($ENV{USE_STATIC_NCCL})
-  message(
-    STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library")
-  set(NCCL_LIBNAME "libnccl_static.a")
-else()
-  set(NCCL_LIBNAME "nccl")
-endif()
-
-find_library(
-  NCCL_LIBRARIES
-  NAMES ${NCCL_LIBNAME}
-  HINTS ${NCCL_LIB_DIR}
-        ${NCCL_ROOT_DIR}
-        ${NCCL_ROOT_DIR}/lib
-        ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu
-        ${NCCL_ROOT_DIR}/lib64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
-
-include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS
-                                  NCCL_LIBRARIES)
-
-if(NCCL_FOUND)
-  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h")
-  message(
-    STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}")
-  file(
-    STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED
-    REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$"
-    LIMIT_COUNT 1)
-  if(NCCL_MAJOR_VERSION_DEFINED)
-    string(REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" ""
-                         NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED})
-    message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}")
-  endif()
-  message(
-    STATUS
-      "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
-  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
-endif()
@@ -1,3 +0,0 @@
-# This file does nothing but to suppress the cmake warning: "By not providing
-# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
-# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.
@@ -1,7 +1,5 @@
 include(CMakeParseArguments)

-# clang format off
-#
 # ##############################################################################
 # Build metal library
 #
@@ -11,14 +9,11 @@ include(CMakeParseArguments)
 # 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) DEBUG: Boolean, if true, enables debug compile options
-# for this specific library. If not provided, uses global MLX_METAL_DEBUG.
+# files (like headers)
 #
-# clang format on
-
 macro(mlx_build_metallib)
  # Parse args
-  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG)
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
  set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
  cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})

@@ -26,11 +21,7 @@ macro(mlx_build_metallib)
  set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")

  # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
-  if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
-    set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
-                               -frecord-sources)
-  endif()
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math)

  # Prepare metallib build command
  add_custom_command(
@@ -13,7 +13,7 @@ EXCLUDE_PATTERNS       = */private/*
 CREATE_SUBDIRS         = NO
 FULL_PATH_NAMES        = YES
 RECURSIVE              = YES
-GENERATE_HTML          = NO
+GENERATE_HTML          = YES
 GENERATE_LATEX         = NO
 GENERATE_XML           = YES
 XML_PROGRAMLISTING     = YES
@@ -26,7 +26,6 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
-PREDEFINED             = MLX_API=

 ################################################################################
 # Compound extraction control.                                                 #
@@ -38,17 +38,3 @@ the docs. Then force add the `build/html` directory:
 `git add -f build/html`

 Commit and push the changes to the `gh-pages` branch.
-
-## Doc Development Setup
-
-To enable live refresh of docs while writing:
-
-Install sphinx autobuild
-```
-pip install sphinx-autobuild
-```
-
-Run auto build on docs/src folder
-```
-sphinx-autobuild ./src ./build/html
-```
@@ -1,36 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
-<defs>
-<g>
-<g id="glyph-0-0">
-<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
-</g>
-<g id="glyph-0-1">
-<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-2">
-<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-3">
-<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
-</g>
-</g>
-<clipPath id="clip-0">
-<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
-</clipPath>
-<clipPath id="clip-1">
-<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
-</clipPath>
-</defs>
-<g clip-path="url(#clip-0)">
-<g fill="rgb(0%, 0%, 0%)" fill-opacity="1">
-<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
-<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
-</g>
-</g>
-<g clip-path="url(#clip-1)">
-<g fill="rgb(82.998657%, 82.998657%, 82.998657%)" fill-opacity="1">
-<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
-</g>
-</g>
-</svg>
@@ -1,36 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
-<defs>
-<g>
-<g id="glyph-0-0">
-<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
-</g>
-<g id="glyph-0-1">
-<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-2">
-<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
-</g>
-<g id="glyph-0-3">
-<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
-</g>
-</g>
-<clipPath id="clip-0">
-<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
-</clipPath>
-<clipPath id="clip-1">
-<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
-</clipPath>
-</defs>
-<g clip-path="url(#clip-0)">
-<g fill="rgb(100%, 100%, 100%)" fill-opacity="1">
-<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
-<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
-</g>
-</g>
-<g clip-path="url(#clip-1)">
-<g fill="rgb(56.999207%, 56.999207%, 56.999207%)" fill-opacity="1">
-<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
-</g>
-</g>
-</svg>
@@ -1,5 +1,4 @@
 sphinx
 breathe
 sphinx-book-theme
-sphinx-copybutton
 mlx
@@ -10,7 +10,7 @@ import mlx.core as mx
 # -- Project information -----------------------------------------------------

 project = "MLX"
-copyright = "2023, Apple"
+copyright = "2023, MLX Contributors"
 author = "MLX Contributors"
 version = ".".join(mx.__version__.split(".")[:3])
 release = version
@@ -18,7 +18,6 @@ release = version
 # -- General configuration ---------------------------------------------------

 extensions = [
-    "sphinx_copybutton",
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.intersphinx",
@@ -8,26 +8,23 @@ MLX supports writing custom Metal kernels through the Python and C++ APIs.
 Simple Example
 --------------

-.. currentmodule:: mlx.core
-
 Let's write a custom kernel that computes ``exp`` elementwise:

 .. code-block:: python

-  source = """
-      uint elem = thread_position_in_grid.x;
-      T tmp = inp[elem];
-      out[elem] = metal::exp(tmp);
-  """
-
-  kernel = mx.fast.metal_kernel(
-      name="myexp",
-      input_names=["inp"],
-      output_names=["out"],
-      source=source,
-  )
-
  def exp_elementwise(a: mx.array):
+      source = """
+          uint elem = thread_position_in_grid.x;
+          T tmp = inp[elem];
+          out[elem] = metal::exp(tmp);
+      """
+
+      kernel = mx.fast.metal_kernel(
+          name="myexp",
+          input_names=["inp"],
+          output_names=["out"],
+          source=source,
+      )
      outputs = kernel(
          inputs=[a],
          template=[("T", mx.float32)],
@@ -42,13 +39,8 @@ Let's write a custom kernel that computes ``exp`` elementwise:
  b = exp_elementwise(a)
  assert mx.allclose(b, mx.exp(a))

-Every time you make a kernel, a new Metal library is created and possibly
-JIT compiled. To reduce the overhead from that, build the kernel once with
-:func:`fast.metal_kernel` and then use it many times.
-
 .. note::
-   Only pass the body of the Metal kernel in ``source``. The function
-   signature is generated automatically.
+    We are only required to pass the body of the Metal kernel in ``source``.

 The full function signature will be generated using:

@@ -86,52 +78,44 @@ 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.
+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 :func:`ast.metal_kernel.__call__` will print the
-generated code for debugging purposes.
+Passing ``verbose=True`` to ``mx.fast.metal_kernel.__call__`` will print the generated code for debugging purposes.

 Using Shape/Strides
 -------------------

-:func:`fast.metal_kernel` supports an argument ``ensure_row_contiguous`` which
-is ``True`` by default. This will copy the array inputs if needed
-before the kernel is launched to ensure that the memory layout is row
-contiguous.  Generally this makes writing the kernel easier, since we don't
-have to worry about gaps or the ordering of the dims when indexing.
+``mx.fast.metal_kernel`` supports an argument ``ensure_row_contiguous`` which is ``True`` by default.
+This will copy the ``mx.array`` inputs if needed before the kernel is launched to ensure that the memory layout is row contiguous.
+Generally this makes writing the kernel easier, since we don't have to worry about gaps or the ordering of the dims
+when indexing.

-If we want to avoid this copy, :func:`fast.metal_kernel` automatically passes
-``a_shape``, ``a_strides`` and ``a_ndim`` for each input array ``a`` if any are
-present in ``source``. We can then use MLX's built in indexing utils to fetch
-the right elements for each thread.
+If we want to avoid this copy, ``metal_kernel`` automatically passes ``a_shape``, ``a_strides`` and ``a_ndim`` for each
+input array ``a`` if any are present in ``source``.
+We can then use MLX's built in indexing utils to fetch the right elements for each thread.

-Let's convert ``myexp`` above to support arbitrarily strided arrays without
-relying on a copy from ``ensure_row_contiguous``:
+Let's convert ``myexp`` above to support arbitrarily strided arrays without relying on a copy from ``ensure_row_contiguous``:

 .. code-block:: python
-   
-  source = """
-      uint elem = thread_position_in_grid.x;
-      // Utils from `mlx/backend/metal/kernels/utils.h` are automatically included
-      uint loc = elem_to_loc(elem, inp_shape, inp_strides, inp_ndim);
-      T tmp = inp[loc];
-      // Output arrays are always row contiguous
-      out[elem] = metal::exp(tmp);
-  """
-
-  kernel = mx.fast.metal_kernel(
-      name="myexp_strided",
-      input_names=["inp"],
-      output_names=["out"],
-      source=source,
-      ensure_row_contiguous=False,
-  )

  def exp_elementwise(a: mx.array):
+      source = """
+          uint elem = thread_position_in_grid.x;
+          // Utils from `mlx/backend/metal/kernels/utils.h` are automatically included
+          uint loc = elem_to_loc(elem, inp_shape, inp_strides, inp_ndim);
+          T tmp = inp[loc];
+          // Output arrays are always row contiguous
+          out[elem] = metal::exp(tmp);
+      """
+
+      kernel = mx.fast.metal_kernel(
+          name="myexp_strided",
+          input_names=["inp"],
+          output_names=["out"],
+          source=source
+      )
      outputs = kernel(
          inputs=[a],
          template=[("T", mx.float32)],
@@ -139,6 +123,7 @@ relying on a copy from ``ensure_row_contiguous``:
          threadgroup=(256, 1, 1),
          output_shapes=[a.shape],
          output_dtypes=[a.dtype],
+          ensure_row_contiguous=False,
      )
      return outputs[0]

@@ -157,139 +142,137 @@ We'll start with the following MLX implementation using standard ops:

 .. code-block:: python

-  def grid_sample_ref(x, grid):
-      N, H_in, W_in, _ = x.shape
-      ix = ((grid[..., 0] + 1) * W_in - 1) / 2
-      iy = ((grid[..., 1] + 1) * H_in - 1) / 2
+    def grid_sample_ref(x, grid):
+        N, H_in, W_in, _ = x.shape
+        ix = ((grid[..., 0] + 1) * W_in - 1) / 2
+        iy = ((grid[..., 1] + 1) * H_in - 1) / 2

-      ix_nw = mx.floor(ix).astype(mx.int32)
-      iy_nw = mx.floor(iy).astype(mx.int32)
+        ix_nw = mx.floor(ix).astype(mx.int32)
+        iy_nw = mx.floor(iy).astype(mx.int32)

-      ix_ne = ix_nw + 1
-      iy_ne = iy_nw
+        ix_ne = ix_nw + 1
+        iy_ne = iy_nw

-      ix_sw = ix_nw
-      iy_sw = iy_nw + 1
+        ix_sw = ix_nw
+        iy_sw = iy_nw + 1

-      ix_se = ix_nw + 1
-      iy_se = iy_nw + 1
+        ix_se = ix_nw + 1
+        iy_se = iy_nw + 1

-      nw = (ix_se - ix)    * (iy_se - iy)
-      ne = (ix    - ix_sw) * (iy_sw - iy)
-      sw = (ix_ne - ix)    * (iy    - iy_ne)
-      se = (ix    - ix_nw) * (iy    - iy_nw)
+        nw = (ix_se - ix)    * (iy_se - iy)
+        ne = (ix    - ix_sw) * (iy_sw - iy)
+        sw = (ix_ne - ix)    * (iy    - iy_ne)
+        se = (ix    - ix_nw) * (iy    - iy_nw)

-      I_nw = x[mx.arange(N)[:, None, None], iy_nw, ix_nw, :]
-      I_ne = x[mx.arange(N)[:, None, None], iy_ne, ix_ne, :]
-      I_sw = x[mx.arange(N)[:, None, None], iy_sw, ix_sw, :]
-      I_se = x[mx.arange(N)[:, None, None], iy_se, ix_se, :]
+        I_nw = x[mx.arange(N)[:, None, None], iy_nw, ix_nw, :]
+        I_ne = x[mx.arange(N)[:, None, None], iy_ne, ix_ne, :]
+        I_sw = x[mx.arange(N)[:, None, None], iy_sw, ix_sw, :]
+        I_se = x[mx.arange(N)[:, None, None], iy_se, ix_se, :]

-      mask_nw = (iy_nw >= 0) & (iy_nw <= H_in - 1) & (ix_nw >= 0) & (ix_nw <= W_in - 1)
-      mask_ne = (iy_ne >= 0) & (iy_ne <= H_in - 1) & (ix_ne >= 0) & (ix_ne <= W_in - 1)
-      mask_sw = (iy_sw >= 0) & (iy_sw <= H_in - 1) & (ix_sw >= 0) & (ix_sw <= W_in - 1)
-      mask_se = (iy_se >= 0) & (iy_se <= H_in - 1) & (ix_se >= 0) & (ix_se <= W_in - 1)
+        mask_nw = (iy_nw >= 0) & (iy_nw <= H_in - 1) & (ix_nw >= 0) & (ix_nw <= W_in - 1)
+        mask_ne = (iy_ne >= 0) & (iy_ne <= H_in - 1) & (ix_ne >= 0) & (ix_ne <= W_in - 1)
+        mask_sw = (iy_sw >= 0) & (iy_sw <= H_in - 1) & (ix_sw >= 0) & (ix_sw <= W_in - 1)
+        mask_se = (iy_se >= 0) & (iy_se <= H_in - 1) & (ix_se >= 0) & (ix_se <= W_in - 1)

-      I_nw *= mask_nw[..., None]
-      I_ne *= mask_ne[..., None]
-      I_sw *= mask_sw[..., None]
-      I_se *= mask_se[..., None]
+        I_nw *= mask_nw[..., None]
+        I_ne *= mask_ne[..., None]
+        I_sw *= mask_sw[..., None]
+        I_se *= mask_se[..., None]

-      output = nw[..., None] * I_nw + ne[..., None] * I_ne + sw[..., None] * I_sw + se[..., None] * I_se
+        output = nw[..., None] * I_nw + ne[..., None] * I_ne + sw[..., None] * I_sw + se[..., None] * I_se

-      return output
+        return output

-Now let's use :func:`custom_function` together with :func:`fast.metal_kernel`
+Now let's use ``mx.custom_function`` together with ``mx.fast.metal_kernel``
 to write a fast GPU kernel for both the forward and backward passes.

 First we'll implement the forward pass as a fused kernel:

 .. code-block:: python

-  source = """
-      uint elem = thread_position_in_grid.x;
-      int H = x_shape[1];
-      int W = x_shape[2];
-      int C = x_shape[3];
-      int gH = grid_shape[1];
-      int gW = grid_shape[2];
+    @mx.custom_function
+    def grid_sample(x, grid):

-      int w_stride = C;
-      int h_stride = W * w_stride;
-      int b_stride = H * h_stride;
+        assert x.ndim == 4, "`x` must be 4D."
+        assert grid.ndim == 4, "`grid` must be 4D."

-      uint grid_idx = elem / C * 2;
-      float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
-      float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;
+        B, _, _, C = x.shape
+        _, gN, gM, D = grid.shape
+        out_shape = (B, gN, gM, C)

-      int ix_nw = floor(ix);
-      int iy_nw = floor(iy);
+        assert D == 2, "Last dim of `grid` must be size 2."

-      int ix_ne = ix_nw + 1;
-      int iy_ne = iy_nw;
+        source = """
+            uint elem = thread_position_in_grid.x;
+            int H = x_shape[1];
+            int W = x_shape[2];
+            int C = x_shape[3];
+            int gH = grid_shape[1];
+            int gW = grid_shape[2];

-      int ix_sw = ix_nw;
-      int iy_sw = iy_nw + 1;
+            int w_stride = C;
+            int h_stride = W * w_stride;
+            int b_stride = H * h_stride;

-      int ix_se = ix_nw + 1;
-      int iy_se = iy_nw + 1;
+            uint grid_idx = elem / C * 2;
+            float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
+            float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;

-      T nw = (ix_se - ix)    * (iy_se - iy);
-      T ne = (ix    - ix_sw) * (iy_sw - iy);
-      T sw = (ix_ne - ix)    * (iy    - iy_ne);
-      T se = (ix    - ix_nw) * (iy    - iy_nw);
+            int ix_nw = floor(ix);
+            int iy_nw = floor(iy);

-      int batch_idx = elem / C / gH / gW * b_stride;
-      int channel_idx = elem % C;
-      int base_idx = batch_idx + channel_idx;
+            int ix_ne = ix_nw + 1;
+            int iy_ne = iy_nw;

-      T I_nw = x[base_idx + iy_nw * h_stride + ix_nw * w_stride];
-      T I_ne = x[base_idx + iy_ne * h_stride + ix_ne * w_stride];
-      T I_sw = x[base_idx + iy_sw * h_stride + ix_sw * w_stride];
-      T I_se = x[base_idx + iy_se * h_stride + ix_se * w_stride];
+            int ix_sw = ix_nw;
+            int iy_sw = iy_nw + 1;

-      I_nw = iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1 ? I_nw : 0;
-      I_ne = iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1 ? I_ne : 0;
-      I_sw = iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1 ? I_sw : 0;
-      I_se = iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1 ? I_se : 0;
+            int ix_se = ix_nw + 1;
+            int iy_se = iy_nw + 1;

-      out[elem] = nw * I_nw + ne * I_ne + sw * I_sw + se * I_se;
-  """
+            T nw = (ix_se - ix)    * (iy_se - iy);
+            T ne = (ix    - ix_sw) * (iy_sw - iy);
+            T sw = (ix_ne - ix)    * (iy    - iy_ne);
+            T se = (ix    - ix_nw) * (iy    - iy_nw);

-  kernel = mx.fast.metal_kernel(
-      name="grid_sample",
-      input_names=["x", "grid"],
-      output_names=["out"],
-      source=source,
-  )
+            int batch_idx = elem / C / gH / gW * b_stride;
+            int channel_idx = elem % C;
+            int base_idx = batch_idx + channel_idx;

-  @mx.custom_function
-  def grid_sample(x, grid):
+            T I_nw = x[base_idx + iy_nw * h_stride + ix_nw * w_stride];
+            T I_ne = x[base_idx + iy_ne * h_stride + ix_ne * w_stride];
+            T I_sw = x[base_idx + iy_sw * h_stride + ix_sw * w_stride];
+            T I_se = x[base_idx + iy_se * h_stride + ix_se * w_stride];

-      assert x.ndim == 4, "`x` must be 4D."
-      assert grid.ndim == 4, "`grid` must be 4D."
+            I_nw = iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1 ? I_nw : 0;
+            I_ne = iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1 ? I_ne : 0;
+            I_sw = iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1 ? I_sw : 0;
+            I_se = iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1 ? I_se : 0;

-      B, _, _, C = x.shape
-      _, gN, gM, D = grid.shape
-      out_shape = (B, gN, gM, C)
-
-      assert D == 2, "Last dim of `grid` must be size 2."
-
-      outputs = kernel(
-          inputs=[x, grid],
-          template=[("T", x.dtype)],
-          output_shapes=[out_shape],
-          output_dtypes=[x.dtype],
-          grid=(np.prod(out_shape), 1, 1),
-          threadgroup=(256, 1, 1),
-      )
-      return outputs[0]
+            out[elem] = nw * I_nw + ne * I_ne + sw * I_sw + se * I_se;
+        """
+        kernel = mx.fast.metal_kernel(
+            name="grid_sample",
+            input_names=["x", "grid"],
+            output_names=["out"],
+            source=source,
+        )
+        outputs = kernel(
+            inputs=[x, grid],
+            template=[("T", x.dtype)],
+            output_shapes=[out_shape],
+            output_dtypes=[x.dtype],
+            grid=(np.prod(out_shape), 1, 1),
+            threadgroup=(256, 1, 1),
+        )
+        return outputs[0]

 For a reasonably sized input such as:

 .. code-block:: python

-  x.shape = (8, 1024, 1024, 64)
-  grid.shape = (8, 256, 256, 2)
+    x.shape = (8, 1024, 1024, 64)
+    grid.shape = (8, 256, 256, 2)

 On an M1 Max, we see a big performance improvement:

@@ -298,11 +281,11 @@ On an M1 Max, we see a big performance improvement:
 Grid Sample VJP
 ---------------

-Since we decorated ``grid_sample`` with :func:`custom_function`, we can now
-define its custom vjp transform so MLX can differentiate it.
+Since we decorated ``grid_sample`` with ``mx.custom_function``, we can now define
+its custom vjp transform so MLX can differentiate it.

 The backwards pass requires atomically updating ``x_grad``/``grid_grad`` and so
-requires a few extra :func:`fast.metal_kernel` features:
+requires a few extra ``mx.fast.metal_kernel`` features:

 * ``init_value=0``
    Initialize all of the kernel's outputs to this value before it runs. This allows us to update only part of the output arrays with the kernel.
@@ -316,129 +299,128 @@ We can then implement the backwards pass as follows:

 .. code-block:: python

-  source = """
-      uint elem = thread_position_in_grid.x;
-      int H = x_shape[1];
-      int W = x_shape[2];
-      int C = x_shape[3];
-      // Pad C to the nearest larger simdgroup size multiple
-      int C_padded = ceildiv(C, threads_per_simdgroup) * threads_per_simdgroup;
+    @grid_sample.vjp
+    def grid_sample_vjp(primals, cotangent, _):
+        x, grid = primals
+        B, _, _, C = x.shape
+        _, gN, gM, D = grid.shape

-      int gH = grid_shape[1];
-      int gW = grid_shape[2];
+        assert D == 2, "Last dim of `grid` must be size 2."

-      int w_stride = C;
-      int h_stride = W * w_stride;
-      int b_stride = H * h_stride;
+        source = """
+            uint elem = thread_position_in_grid.x;
+            int H = x_shape[1];
+            int W = x_shape[2];
+            int C = x_shape[3];
+            // Pad C to the nearest larger simdgroup size multiple
+            int C_padded = ceildiv(C, threads_per_simdgroup) * threads_per_simdgroup;

-      uint grid_idx = elem / C_padded * 2;
-      float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
-      float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;
+            int gH = grid_shape[1];
+            int gW = grid_shape[2];

-      int ix_nw = floor(ix);
-      int iy_nw = floor(iy);
+            int w_stride = C;
+            int h_stride = W * w_stride;
+            int b_stride = H * h_stride;

-      int ix_ne = ix_nw + 1;
-      int iy_ne = iy_nw;
+            uint grid_idx = elem / C_padded * 2;
+            float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
+            float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;

-      int ix_sw = ix_nw;
-      int iy_sw = iy_nw + 1;
+            int ix_nw = floor(ix);
+            int iy_nw = floor(iy);

-      int ix_se = ix_nw + 1;
-      int iy_se = iy_nw + 1;
+            int ix_ne = ix_nw + 1;
+            int iy_ne = iy_nw;

-      T nw = (ix_se - ix)    * (iy_se - iy);
-      T ne = (ix    - ix_sw) * (iy_sw - iy);
-      T sw = (ix_ne - ix)    * (iy    - iy_ne);
-      T se = (ix    - ix_nw) * (iy    - iy_nw);
+            int ix_sw = ix_nw;
+            int iy_sw = iy_nw + 1;

-      int batch_idx = elem / C_padded / gH / gW * b_stride;
-      int channel_idx = elem % C_padded;
-      int base_idx = batch_idx + channel_idx;
+            int ix_se = ix_nw + 1;
+            int iy_se = iy_nw + 1;

-      T gix = T(0);
-      T giy = T(0);
-      if (channel_idx < C) {
-          int cot_index = elem / C_padded * C + channel_idx;
-          T cot = cotangent[cot_index];
-          if (iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1) {
-              int offset = base_idx + iy_nw * h_stride + ix_nw * w_stride;
-              atomic_fetch_add_explicit(&x_grad[offset], nw * cot, memory_order_relaxed);
+            T nw = (ix_se - ix)    * (iy_se - iy);
+            T ne = (ix    - ix_sw) * (iy_sw - iy);
+            T sw = (ix_ne - ix)    * (iy    - iy_ne);
+            T se = (ix    - ix_nw) * (iy    - iy_nw);

-              T I_nw = x[offset];
-              gix -= I_nw * (iy_se - iy) * cot;
-              giy -= I_nw * (ix_se - ix) * cot;
-          }
-          if (iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1) {
-              int offset = base_idx + iy_ne * h_stride + ix_ne * w_stride;
-              atomic_fetch_add_explicit(&x_grad[offset], ne * cot, memory_order_relaxed);
+            int batch_idx = elem / C_padded / gH / gW * b_stride;
+            int channel_idx = elem % C_padded;
+            int base_idx = batch_idx + channel_idx;

-              T I_ne = x[offset];
-              gix += I_ne * (iy_sw - iy) * cot;
-              giy -= I_ne * (ix - ix_sw) * cot;
-          }
-          if (iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1) {
-              int offset = base_idx + iy_sw * h_stride + ix_sw * w_stride;
-              atomic_fetch_add_explicit(&x_grad[offset], sw * cot, memory_order_relaxed);
+            T gix = T(0);
+            T giy = T(0);
+            if (channel_idx < C) {
+                int cot_index = elem / C_padded * C + channel_idx;
+                T cot = cotangent[cot_index];
+                if (iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1) {
+                    int offset = base_idx + iy_nw * h_stride + ix_nw * w_stride;
+                    atomic_fetch_add_explicit(&x_grad[offset], nw * cot, memory_order_relaxed);

-              T I_sw = x[offset];
-              gix -= I_sw * (iy - iy_ne) * cot;
-              giy += I_sw * (ix_ne - ix) * cot;
-          }
-          if (iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1) {
-              int offset = base_idx + iy_se * h_stride + ix_se * w_stride;
-              atomic_fetch_add_explicit(&x_grad[offset], se * cot, memory_order_relaxed);
+                    T I_nw = x[offset];
+                    gix -= I_nw * (iy_se - iy) * cot;
+                    giy -= I_nw * (ix_se - ix) * cot;
+                }
+                if (iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1) {
+                    int offset = base_idx + iy_ne * h_stride + ix_ne * w_stride;
+                    atomic_fetch_add_explicit(&x_grad[offset], ne * cot, memory_order_relaxed);

-              T I_se = x[offset];
-              gix += I_se * (iy - iy_nw) * cot;
-              giy += I_se * (ix - ix_nw) * cot;
-          }
-      }
+                    T I_ne = x[offset];
+                    gix += I_ne * (iy_sw - iy) * cot;
+                    giy -= I_ne * (ix - ix_sw) * cot;
+                }
+                if (iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1) {
+                    int offset = base_idx + iy_sw * h_stride + ix_sw * w_stride;
+                    atomic_fetch_add_explicit(&x_grad[offset], sw * cot, memory_order_relaxed);

-      T gix_mult = W / 2;
-      T giy_mult = H / 2;
+                    T I_sw = x[offset];
+                    gix -= I_sw * (iy - iy_ne) * cot;
+                    giy += I_sw * (ix_ne - ix) * cot;
+                }
+                if (iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1) {
+                    int offset = base_idx + iy_se * h_stride + ix_se * w_stride;
+                    atomic_fetch_add_explicit(&x_grad[offset], se * cot, memory_order_relaxed);

-      // Reduce across each simdgroup first.
-      // This is much faster than relying purely on atomics.
-      gix = simd_sum(gix);
-      giy = simd_sum(giy);
+                    T I_se = x[offset];
+                    gix += I_se * (iy - iy_nw) * cot;
+                    giy += I_se * (ix - ix_nw) * cot;
+                }
+            }

-      if (thread_index_in_simdgroup == 0) {
-          atomic_fetch_add_explicit(&grid_grad[grid_idx], gix * gix_mult, memory_order_relaxed);
-          atomic_fetch_add_explicit(&grid_grad[grid_idx + 1], giy * giy_mult, memory_order_relaxed);
-      }
-  """
-  kernel = mx.fast.metal_kernel(
-      name="grid_sample_grad",
-      input_names=["x", "grid", "cotangent"],
-      output_names=["x_grad", "grid_grad"],
-      source=source,
-      atomic_outputs=True,
-  )
+            T gix_mult = W / 2;
+            T giy_mult = H / 2;

-  @grid_sample.vjp
-  def grid_sample_vjp(primals, cotangent, _):
-      x, grid = primals
-      B, _, _, C = x.shape
-      _, gN, gM, D = grid.shape
+            // Reduce across each simdgroup first.
+            // This is much faster than relying purely on atomics.
+            gix = simd_sum(gix);
+            giy = simd_sum(giy);

-      assert D == 2, "Last dim of `grid` must be size 2."
-
-      # pad the output channels to simd group size
-      # so that our `simd_sum`s don't overlap.
-      simdgroup_size = 32
-      C_padded = (C + simdgroup_size - 1) // simdgroup_size * simdgroup_size
-      grid_size = B * gN * gM * C_padded
-      outputs = kernel(
-          inputs=[x, grid, cotangent],
-          template=[("T", x.dtype)],
-          output_shapes=[x.shape, grid.shape],
-          output_dtypes=[x.dtype, x.dtype],
-          grid=(grid_size, 1, 1),
-          threadgroup=(256, 1, 1),
-          init_value=0,
-      )
-      return outputs[0], outputs[1]
+            if (thread_index_in_simdgroup == 0) {
+                atomic_fetch_add_explicit(&grid_grad[grid_idx], gix * gix_mult, memory_order_relaxed);
+                atomic_fetch_add_explicit(&grid_grad[grid_idx + 1], giy * giy_mult, memory_order_relaxed);
+            }
+        """
+        kernel = mx.fast.metal_kernel(
+            name="grid_sample_grad",
+            input_names=["x", "grid", "cotangent"],
+            output_names=["x_grad", "grid_grad"],
+            source=source,
+            atomic_outputs=True,
+        )
+        # pad the output channels to simd group size
+        # so that our `simd_sum`s don't overlap.
+        simdgroup_size = 32
+        C_padded = (C + simdgroup_size - 1) // simdgroup_size * simdgroup_size
+        grid_size = B * gN * gM * C_padded
+        outputs = kernel(
+            inputs=[x, grid, cotangent],
+            template=[("T", x.dtype)],
+            output_shapes=[x.shape, grid.shape],
+            output_dtypes=[x.dtype, x.dtype],
+            grid=(grid_size, 1, 1),
+            threadgroup=(256, 1, 1),
+            init_value=0,
+        )
+        return outputs[0], outputs[1]

 There's an even larger speed up for the vjp:

@@ -22,12 +22,12 @@ You can do that in MLX directly:
 This function performs that operation while leaving the implementation and
 function transformations to MLX.

-However, you may want to customize the underlying implementation, perhaps to
-make it faster. In this tutorial we will go through adding custom extensions.
-It will cover:
+However you may need to customize the underlying implementation, perhaps to
+make it faster or for custom differentiation. In this tutorial we will go
+through adding custom extensions. It will cover:

 * The structure of the MLX library.
-* Implementing a CPU operation.
+* Implementing a CPU operation that redirects to Accelerate_ when appropriate.
 * Implementing a GPU operation using metal.
 * Adding the ``vjp`` and ``jvp`` function transformation.
 * Building a custom extension and binding it to python.
@@ -45,7 +45,7 @@ Operations
 Operations are the front-end functions that operate on arrays. They are defined
 in the C++ API (:ref:`cpp_ops`), and the Python API (:ref:`ops`) binds them.

-We would like an operation :meth:`axpby` that takes in two arrays, ``x`` and
+We would like an operation, :meth:`axpby` that takes in two arrays ``x`` and
 ``y``, and two scalars, ``alpha`` and ``beta``. This is how to define it in
 C++:

@@ -55,7 +55,7 @@ C++:
    *  Scale and sum two vectors element-wise
    *  z = alpha * x + beta * y
    *
-    *  Use NumPy-style broadcasting between x and y
+    *  Follow numpy style broadcasting between x and y
    *  Inputs are upcasted to floats if needed
    **/
    array axpby(
@@ -66,7 +66,7 @@ C++:
        StreamOrDevice s = {} // Stream on which to schedule the operation
    );

-The simplest way to implement this is with existing operations:
+The simplest way to this operation is in terms of existing operations:

 .. code-block:: C++

@@ -93,9 +93,9 @@ Primitives
 ^^^^^^^^^^^

 A :class:`Primitive` is part of the computation graph of an :class:`array`. It
-defines how to create output arrays given input arrays. Further, a
+defines how to create outputs arrays given a input arrays. Further, a
 :class:`Primitive` has methods to run on the CPU or GPU and for function
-transformations such as ``vjp`` and ``jvp``.  Let's go back to our example to be
+transformations such as ``vjp`` and ``jvp``.  Lets go back to our example to be
 more concrete:

 .. code-block:: C++
@@ -128,7 +128,7 @@ more concrete:
        /** The vector-Jacobian product. */
        std::vector<array> vjp(
            const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
+            const array& cotan,
            const std::vector<int>& argnums,
            const std::vector<array>& outputs) override;

@@ -138,13 +138,13 @@ more concrete:
        * representing the vectorized computation and the axis which
        * corresponds to the output vectorized dimension.
        */
-        std::pair<std::vector<array>, std::vector<int>> vmap(
+        virtual std::pair<std::vector<array>, std::vector<int>> vmap(
            const std::vector<array>& inputs,
            const std::vector<int>& axes) override;

-        /** The name of primitive. */
-        const char* name() const override {
-          return "Axpby";
+        /** Print the primitive. */
+        void print(std::ostream& os) override {
+            os << "Axpby";
        }

        /** Equivalence check **/
@@ -153,6 +153,9 @@ more concrete:
      private:
        float alpha_;
        float beta_;
+
+        /** Fall back implementation for evaluation on CPU */
+        void eval(const std::vector<array>& inputs, array& out);
    };

 The :class:`Axpby` class derives from the base :class:`Primitive` class. The
@@ -185,7 +188,7 @@ Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
        auto promoted_dtype = promote_types(x.dtype(), y.dtype());

        // Upcast to float32 for non-floating point inputs x and y
-        auto out_dtype = issubdtype(promoted_dtype, float32)
+        auto out_dtype = is_floating_point(promoted_dtype)
            ? promoted_dtype
            : promote_types(promoted_dtype, float32);

@@ -231,57 +234,49 @@ the execution of the computation graph, and calls :meth:`Axpby::eval_cpu` or
 Implementing the CPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

-Let's start by implementing :meth:`Axpby::eval_cpu`.
+Let's start by implementing a naive and generic version of
+:meth:`Axpby::eval_cpu`. We declared this as a private member function of
+:class:`Axpby` earlier called :meth:`Axpby::eval`.

-The method will go over each element of the output array, find the
+Our naive method will go over each element of the output array, find the
 corresponding input elements of ``x`` and ``y`` and perform the operation
 point-wise. This is captured in the templated function :meth:`axpby_impl`.

 .. code-block:: C++

-  template <typename T>
-  void axpby_impl(
-      const mx::array& x,
-      const mx::array& y,
-      mx::array& out,
-      float alpha_,
-      float beta_,
-      mx::Stream stream) {
-    out.set_data(mx::allocator::malloc(out.nbytes()));
+    template <typename T>
+    void axpby_impl(
+            const array& x,
+            const array& y,
+            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()));

-    // Get the CPU command encoder and register input and output arrays
-    auto& encoder = mx::cpu::get_command_encoder(stream);
-    encoder.set_input_array(x);
-    encoder.set_input_array(y);
-    encoder.set_output_array(out);
+        // Collect input and output data pointers
+        const T* x_ptr = x.data<T>();
+        const T* y_ptr = y.data<T>();
+        T* out_ptr = out.data<T>();

-    // Launch the CPU kernel
-    encoder.dispatch([x_ptr = x.data<T>(),
-                      y_ptr = y.data<T>(),
-                      out_ptr = out.data<T>(),
-                      size = out.size(),
-                      shape = out.shape(),
-                      x_strides = x.strides(),
-                      y_strides = y.strides(),
-                      alpha_,
-                      beta_]() {
+        // Cast alpha and beta to the relevant types
+        T alpha = static_cast<T>(alpha_);
+        T beta = static_cast<T>(beta_);

-      // Cast alpha and beta to the relevant types
-      T alpha = static_cast<T>(alpha_);
-      T beta = static_cast<T>(beta_);
+        // 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());

-      // Do the element-wise operation for each output
-      for (size_t out_idx = 0; out_idx < size; out_idx++) {
-        // Map linear indices to offsets in x and y
-        auto x_offset = mx::elem_to_loc(out_idx, shape, x_strides);
-        auto y_offset = mx::elem_to_loc(out_idx, 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
-        out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
-      }
-    });
-  }
+            // We allocate the output to be contiguous and regularly strided
+            // (defaults to row major) and hence it doesn't need additional mapping
+            out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
+        }
+    }

 Our implementation should work for all incoming floating point arrays.
 Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
@@ -289,32 +284,112 @@ Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and

 .. code-block:: C++

-    void Axpby::eval_cpu(
-        const std::vector<mx::array>& inputs,
-        std::vector<mx::array>& outputs) {
-      auto& x = inputs[0];
-      auto& y = inputs[1];
-      auto& out = outputs[0];
+    /** Fall back implementation for evaluation on CPU */
+    void Axpby::eval(
+      const std::vector<array>& inputs,
+      const std::vector<array>& outputs) {
+        auto& x = inputs[0];
+        auto& y = inputs[1];
+        auto& out = outputs[0];

-      // Dispatch to the correct dtype
-      if (out.dtype() == mx::float32) {
-        return axpby_impl<float>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::float16) {
-        return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::bfloat16) {
-        return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::complex64) {
-        return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_, stream());
-      } else {
-        throw std::runtime_error(
-            "Axpby is only supported for floating point types.");
-      }
+        // Dispatch to the correct dtype
+        if (out.dtype() == 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 {
+            throw std::runtime_error(
+                "[Axpby] Only supports floating point types.");
+        }
+    }
+
+This is good as a fallback implementation. We can use the ``axpby`` routine
+provided by the Accelerate_ framework for a faster implementation in certain
+cases:
+
+#.  Accelerate does not provide implementations of ``axpby`` for half precision
+    floats. We can only use it for ``float32`` types.
+#.  Accelerate assumes the inputs ``x`` and ``y`` are contiguous and all
+    elements have fixed strides between them. We only direct to Accelerate
+    if both ``x`` and ``y`` are row contiguous or column contiguous.
+#.  Accelerate performs the routine ``Y = (alpha * X) + (beta * Y)`` in-place.
+    MLX expects to write the output to a new array. We must copy the elements
+    of ``y`` into the output and use that as an input to ``axpby``.
+
+Let's write an implementation that uses Accelerate in the right conditions.
+It allocates data for the output, copies ``y`` into it, and then calls the
+:func:`catlas_saxpby` from accelerate.
+
+.. code-block:: C++
+
+    template <typename T>
+    void axpby_impl_accelerate(
+            const array& x,
+            const array& y,
+            array& out,
+            float alpha_,
+            float beta_) {
+        // Accelerate library provides catlas_saxpby which does
+        // Y = (alpha * X) + (beta * Y) in place
+        // To use it, we first copy the data in y over to the output array
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+
+        // We then copy over the elements using the contiguous vector specialization
+        copy_inplace(y, out, CopyType::Vector);
+
+        // Get x and y pointers for catlas_saxpby
+        const T* x_ptr = x.data<T>();
+        T* y_ptr = out.data<T>();
+
+        T alpha = static_cast<T>(alpha_);
+        T beta = static_cast<T>(beta_);
+
+        // Call the inplace accelerate operator
+        catlas_saxpby(
+            /* N = */ out.size(),
+            /* ALPHA = */ alpha,
+            /* X = */ x_ptr,
+            /* INCX = */ 1,
+            /* BETA = */ beta,
+            /* Y = */ y_ptr,
+            /* INCY = */ 1);
+    }
+
+For inputs that do not fit the criteria for accelerate, we fall back to
+:meth:`Axpby::eval`. With this in mind, let's finish our
+:meth:`Axpby::eval_cpu`.
+
+.. code-block:: C++
+
+    /** Evaluate primitive on CPU using accelerate specializations */
+    void Axpby::eval_cpu(
+      const std::vector<array>& inputs,
+      const std::vector<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 &&
+            ((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_);
+            return;
+        }
+
+        // Fall back to common back-end if specializations are not available
+        eval(inputs, outputs);
    }

 Just this much is enough to run the operation :meth:`axpby` on a CPU stream! If
 you do not plan on running the operation on the GPU or using transforms on
 computation graphs that contain :class:`Axpby`, you can stop implementing the
-primitive here.
+primitive here and enjoy the speed-ups you get from the Accelerate library.

 Implementing the GPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -345,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 int64_t* x_strides [[buffer(6)]],
-            constant const int64_t* y_strides [[buffer(7)]],
+            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]]) {
        // Convert linear indices to offsets in array
@@ -363,10 +438,24 @@ each instantiation a unique host name so we can identify it.

 .. code-block:: C++

-    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)
+    #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);

 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
@@ -391,20 +480,20 @@ below.
        auto& d = metal::device(s.device);

        // Allocate output memory
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));

        // Resolve name of kernel
-        std::stream kname;
-        kname = "axpby_general_" + type_to_name(out);
+        std::ostringstream kname;
+        kname << "axpby_" << "general_" << type_to_name(out);

-        // Load the metal library
-        auto lib = d.get_library("mlx_ext", current_binary_dir());
+        // Make sure the metal library is available
+        d.register_library("mlx_ext");

        // Make a kernel from this metal library
-        auto kernel = d.get_kernel(kname, lib);
+        auto kernel = d.get_kernel(kname.str(), "mlx_ext");

        // Prepare to encode kernel
-        auto& compute_encoder = mx::metal::get_command_encoder(s);
+        auto& compute_encoder = d.get_command_encoder(s.index);
        compute_encoder.set_compute_pipeline_state(kernel);

        // Kernel parameters are registered with buffer indices corresponding to
@@ -448,7 +537,7 @@ We can now call the :meth:`axpby` operation on both the CPU and the GPU!

 A few things to note about MLX and Metal before moving on. MLX keeps track of
 the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is
-associated. We rely on :meth:`metal::get_command_encoder` to give us the active
+associated. We rely on :meth:`d.get_command_encoder` to give us the active
 metal compute command encoder instead of building a new one and calling
 :meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute
 pipelines) to the active command buffer until some specified limit is hit or
@@ -469,7 +558,7 @@ one we just defined:
            const std::vector<array>& tangents,
            const std::vector<int>& argnums) {
        // Forward mode diff that pushes along the tangents
-        // The jvp transform on the primitive can be built with ops
+        // The jvp transform on the primitive can built with ops
        // that are scheduled on the same stream as the primitive

        // If argnums = {0}, we only push along x in which case the
@@ -481,7 +570,7 @@ one we just defined:
            auto scale_arr = array(scale, tangents[0].dtype());
            return {multiply(scale_arr, tangents[0], stream())};
        }
-        // If argnums = {0, 1}, we take contributions from both
+        // If, argnums = {0, 1}, we take contributions from both
        // which gives us jvp = tangent_x * alpha + tangent_y * beta
        else {
            return {axpby(tangents[0], tangents[1], alpha_, beta_, stream())};
@@ -735,7 +824,7 @@ Let's look at a simple script and its results:

    print(f"c shape: {c.shape}")
    print(f"c dtype: {c.dtype}")
-    print(f"c is correct: {mx.all(c == 6.0).item()}")
+    print(f"c correct: {mx.all(c == 6.0).item()}")

 Output:

@@ -743,13 +832,13 @@ Output:

    c shape: [3, 4]
    c dtype: float32
-    c is correct: True
+    c correctness: True

 Results
 ^^^^^^^

 Let's run a quick benchmark and see how our new ``axpby`` operation compares
-with the naive :meth:`simple_axpby` we first defined.
+with the naive :meth:`simple_axpby` we first defined on the CPU.

 .. code-block:: python

@@ -757,11 +846,13 @@ with the naive :meth:`simple_axpby` we first defined.
    from mlx_sample_extensions import axpby
    import time

+    mx.set_default_device(mx.cpu)
+
    def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
        return alpha * x + beta * y

-    M = 4096
-    N = 4096
+    M = 256
+    N = 512

    x = mx.random.normal((M, N))
    y = mx.random.normal((M, N))
@@ -772,24 +863,24 @@ with the naive :meth:`simple_axpby` we first defined.

    def bench(f):
        # Warm up
-        for i in range(5):
+        for i in range(100):
            z = f(x, y, alpha, beta)
            mx.eval(z)

        # Timed run
-        s = time.perf_counter()
-        for i in range(100):
+        s = time.time()
+        for i in range(5000):
            z = f(x, y, alpha, beta)
            mx.eval(z)
-        e = time.perf_counter()
-        return 1000 * (e - s) / 100
+        e = time.time()
+        return e - s

    simple_time = bench(simple_axpby)
    custom_time = bench(axpby)

-    print(f"Simple axpby: {simple_time:.3f} ms | Custom axpby: {custom_time:.3f} ms")
+    print(f"Simple axpby: {simple_time:.3f} s | Custom axpby: {custom_time:.3f} s")

-The results are ``Simple axpby: 1.559 ms | Custom axpby: 0.774 ms``. We see
+The results are ``Simple axpby: 0.114 s | Custom axpby: 0.109 s``. We see
 modest improvements right away!

 This operation is now good to be used to build other operations, in
@@ -1,40 +0,0 @@
-Metal Logging
-=============
-
-In debug builds, MLX compiles Metal kernels with ``os_log`` enabled so shader
-warnings and debug messages are visible during development.
-
-.. note::
-    Metal logging is only available with Metal 3.2 or higher (macOS 15 and up,
-    iOS 18 and up).
-
-To enable logging from kernels, first make sure to build in debug mode:
-
-.. code-block:: bash
-
-    DEBUG=1 python -m pip install -e .
-
-Then, in the kernel source code include MLX's logging shim and use
-``mlx::os_log``:
-
-.. code-block::
-
-    #include "mlx/backend/metal/kernels/logging.h"
-
-    constant mlx::os_log logger("mlx", "my_kernel");
-
-    kernel void my_kernel(/* ... */) {
-    // ...
-      logger.log_debug("unexpected state: idx=%u", idx);
-    }
-
-When you run the program, set the Metal log level to your desired level and
-forward logs to ``stderr``:
-
-.. code-block:: bash
-
-    MTL_LOG_LEVEL=MTLLogLevelDebug MTL_LOG_TO_STDERR=1 python script.py
-
-See the `Metal logging guide`_ for more details.
-
-.. _`Metal logging guide`: https://developer.apple.com/documentation/metal/logging-shader-debug-messages
@@ -1,121 +0,0 @@
-.. _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 20)
-  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
@@ -1,91 +0,0 @@
-.. _data_parallelism:
-
-Data Parallelism
-================
-
-MLX enables efficient data parallel distributed training through its
-distributed communication primitives.
-
-.. _training_example:
-
-Training Example
----------------
-
-In this section we will adapt an MLX training loop to support data parallel
-distributed training. Namely, we will average the gradients across a set of
-hosts before applying them to the model.
-
-Our training loop looks like the following code snippet if we omit the model,
-dataset, and optimizer initialization.
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())
-
-All we have to do to average the gradients across machines is perform an
-:func:`all_sum` and divide by the size of the :class:`Group`. Namely we
-have to :func:`mlx.utils.tree_map` the gradients with following function.
-
-.. code:: python
-
-    def all_avg(x):
-        return mx.distributed.all_sum(x) / mx.distributed.init().size()
-
-Putting everything together our training loop step looks as follows with
-everything else remaining the same.
-
-.. code:: python
-
-    from mlx.utils import tree_map
-
-    def all_reduce_grads(grads):
-        N = mx.distributed.init().size()
-        if N == 1:
-            return grads
-        return tree_map(
-            lambda x: mx.distributed.all_sum(x) / N,
-            grads
-        )
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = all_reduce_grads(grads)  # <--- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-Using ``nn.average_gradients``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Although the code example above works correctly; it performs one communication
-per gradient. It is significantly more efficient to aggregate several gradients
-together and perform fewer communication steps.
-
-This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
-almost identical to the example above:
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = mx.nn.average_gradients(grads)  # <---- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())
@@ -1,239 +0,0 @@
-.. _tensor_parallelism:
-
-Tensor Parallelism
-==================
-
-In this example, we will explore how tensor parallelism (TP) works in MLX.  We
-will start with an overview of the distributed layers in ``mlx.nn`` and then
-show how to do tensor parallelism Llama-style transformer models.
-
-Sharded Layers
--------------
-
-:class:`AllToShardedLinear <mlx.nn.AllToShardedLinear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This layer replicates a common input and shards the weight matrix along the
-output dimension across all devices in the :class:`mlx.core.distributed.Group`.
-The layer produces a sharded output.
-
-For example, consider an :class:`mlx.nn.AllToShardedLinear` layer with
-``input_dims=2`` and ``output_dims=2``, a batched input of shape ``(4, 2)``,
-and a device group with 2 devices. The layer shards the weight matrix along the
-output dimension across the two devices, where each device receives the full
-input and computes a partial output.
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/all-to-sharded-linear.png" alt="column-wise tensor parallelism" style="width: 100%">
-    </div>
-
-This layer does not automatically gather all outputs from each device. This is
-an intended and :ref:`useful design choice <useful_design_choices>`.
-
-:class:`QuantizedAllToShardedLinear <mlx.nn.QuantizedAllToShardedLinear>` is
-the quantized equivalent of :class:`mlx.nn.AllToShardedLinear`.  Similar to
-:class:`mlx.nn.QuantizedLinear`, its parameters are frozen and will not be
-included in any gradient computation.
-
-
-:class:`ShardedToAllLinear <mlx.nn.ShardedToAllLinear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This layer expects inputs that are sharded along the feature dimension and
-shards the weight matrix along the input dimension across all devices in the
-:class:`mlx.core.distributed.Group`. The layer automatically aggregates the
-results using :class:`mlx.core.distributed.all_sum`, so all devices in the
-group will have the same result.
-
-For example, consider an :class:`mlx.nn.ShardedToAllLinear` layer with
-``input_dims=2`` and ``output_dims=2``, a batched input of shape ``(4, 2)``,
-and a device group with 2 devices. The layer shards the weight matrix along the
-input dimension across the two devices. Each device computes a ``(4,2)``
-output, which is then aggregated with all other device outputs to get layer
-output.
-
-   .. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/sharded-to-all-linear.png" alt="row-wise tensor parallelism" style="width: 100%">
-    </div>
-
-This layer does not automatically shard the inputs along the feature dimension
-for you. It is necessary to create a "partial" input structure to feed into the
-layer. This is an intended and :ref:`useful design choice
-<useful_design_choices>`.
-
-:class:`QuantizedShardedToAllLinear <mlx.nn.QuantizedShardedToAllLinear>` is
-the quantized equivalent of :class:`mlx.nn.ShardedToAllLinear`.  Similar to
-:class:`mlx.nn.QuantizedLinear`, its parameters are frozen and will not be
-included in any gradient computation.
-
-
-Shard Utility Functions
-----------------------
-
-:func:`shard_linear <mlx.nn.layers.distributed.shard_linear>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Converts a regular linear layer into a tensor parallel layer that distributes
-computation across multiple devices. Takes an existing :class:`mlx.nn.Linear`
-or :class:`mlx.nn.QuantizedLinear` layer and returns a new distributed layer
-(either :class:`mlx.nn.AllToShardedLinear` or
-:class:`mlx.nn.ShardedToAllLinear`, depending on the sharding type). The
-original layer is not modified.
-
-:func:`shard_inplace <mlx.nn.layers.distributed.shard_inplace>`
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Splits the parameters of an existing layer across multiple devices by modifying
-the layer in-place. Unlike :func:`shard_linear
-<mlx.nn.layers.distributed.shard_linear>`, this function does not create a new
-layer or add distributed communication. The layer itself must handle
-distributed communication if needed.
-
-
-.. _useful_design_choices:
-
-Useful Design Choices
---------------------
-
-The design choices above regarding when operations are done automatically are intentional and make model training and inference easier.
-
-All-to-sharded and sharded-to-all layers naturally go together because the
-output of the former layer is exactly the input needed needed for the latter.
-This removes the need for an intermediate gather step between the layers,
-reducing communication overhead.
-
-This is why :class:`mlx.nn.AllToShardedLinear` does not aggregate results
-automatically and why :class:`mlx.nn.ShardedToAllLinear` does not shard inputs
-automatically. It is so that they can be placed in successive order and work
-together easily.
-
-We can demonstrate this through a simple model using our two types of
-distributed layers.
-
-.. code-block:: python
-
-  x = ... # some (4, 2) model input: batch size 4, feature size 2
-
-  l1 = nn.AllToShardedLinear(2, 2, bias=False)   # initialize the layer
-  l1_out = l1(x) # (4, 1) output
-
-  l2 = nn.ShardedToAllLinear(2, 2, bias=False)
-  l2_out = l2(l1_out) # (4, 2) output
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/column-row-tp.png" alt="two layer tensor parallelism" style="width: 100%">
-      <p style="font-size: 0.85em; margin-top: 0.5em;"><small>A visualization of the simple MLX model using all-to-sharded then sharded-to-all tensor parallelism across 2 devices.</small></p>
-    </div>
-
-
-LLM Inference with Tensor Parallelism
-------------------------------------
-
-We can apply these TP techniques to LLMs in order to enable inference for much
-larger models by sharding parameters from huge layers across multiple devices.
-
-To demonstrate this, let's apply TP to the Transformer block of our :doc:`Llama
-Inference <llama-inference>` example. In this example, we will use the same
-inference script as the Llama Inference example, which can be found in
-`mlx-examples`_.
-
-Our first edit is to initialize the distributed communication group and get the
-current process rank:
-
-.. code-block:: python
-
-  world = mx.distributed.init()
-  rank = world.rank()
-
-Next, let's look at the current architecture of the transformer block and see how we can apply tensor parallelism:
-
-.. raw:: html
-
-    <div>
-      <img src="../_static/tp_inference/llama-transformer.png" alt="llama transformer example" style="width: 100%">
-    </div>
-
-
-This architecture has two natural places where 
-tensor parallelism can be applied: the attention block and the FFN
-block. Both follow the same pattern: multiple parallel linear layers operating
-on the same input, followed by a single output linear layer. In the attention
-block, the Q, K, and V projections are sharded along the output dimension (all-to-sharded), and the output
-projection is sharded along the input dimension (sharded-to-all). Similarly in the FFN block, the gate and up projections
-become all-to-sharded layers, and the down projection becomes an sharded-to-all layer.
-
-The intermediate operations between the linear layers (RoPE, softmax, scaled
-dot-product attention in the attention block, and element-wise multiplication
-in the FFN block) do not impede the use of our TP paradigm. These operations
-are either:
-
- **Element-wise operations** (RoPE, element-wise multiplication): These
-  operate independently on each element or position, preserving the sharding
-  pattern without requiring cross-device communication.
-
- **Operations on non-sharded dimensions** (softmax, scaled dot-product
-  attention): These operate along dimensions that are not sharded (such as the
-  sequence length or head dimensions), so they can be computed independently on
-  each device. The attention computation ``Q @ K^T`` and ``scores @ V`` work
-  correctly with sharded Q, K, V tensors because the matrix multiplications are
-  performed along the sharded feature dimension, and the results remain
-  properly sharded for the subsequent sharded-to-all layer.
-
-To implement sharding in our Llama inference, we use :func:`shard_linear
-<mlx.nn.layers.distributed.shard_linear>` to get sharded linear layers with
-distributed communication. This is easier than using :func:`shard_inplace
-<mlx.nn.layers.distributed.shard_inplace>` and implementing the steps manually
-in the :code:`__call__` function.
-
-The following code shows how to shard the Attention block. The Q, K, and V
-projection layers are converted to all-to-sharded layers, while the output
-projection is converted to a sharded-to-all layer. The number of heads are also
-adjusted to account for the sharding:
-
-.. code-block:: python
-
-  # ... in Attention class
-  def shard(self, group: mx.distributed.Group):
-    self.n_heads = self.n_heads // group.size()
-    self.n_kv_heads = self.n_kv_heads // group.size()
-
-    self.wq = nn.layers.distributed.shard_linear(self.wq, "all-to-sharded", group=group)
-    self.wk = nn.layers.distributed.shard_linear(self.wk, "all-to-sharded", group=group)
-    self.wv = nn.layers.distributed.shard_linear(self.wv, "all-to-sharded", group=group)
-    self.wo = nn.layers.distributed.shard_linear(self.wo, "sharded-to-all", group=group)
-
-Similarly, the FeedForward block is sharded by converting the gate (w1) and up
-(w3) projections to all-to-sharded layers, and the down projection (w2) to
-a sharded-to-all layer:
-
-.. code-block:: python
-
-  # ... in FeedForward class
-  def shard(self, group: mx.distributed.Group):
-    self.w1 = nn.layers.distributed.shard_linear(self.w1, "all-to-sharded", group=group)
-    self.w2 = nn.layers.distributed.shard_linear(self.w2, "sharded-to-all", group=group)
-    self.w3 = nn.layers.distributed.shard_linear(self.w3, "all-to-sharded", group=group)
-
-Finally, in our :code:`load_model` function, we need to apply our sharding
-functions to all transformer layers when using multiple devices:
-
-.. code-block:: python
-
-  # ... in load_model function
-  if world.size() > 1:
-    # convert Linear layers in Transformer/FFN to appropriate Sharded Layers
-    for layer in model.layers:
-        layer.attention.shard(group=world)
-        layer.feed_forward.shard(group=world)
-
-This allows us to use the llama inference file as normal when running
-:code:`python llama.py`, but now we can also run it across two (or more)
-devices via :code:`mlx.launch -n 2 llama.py`.
-
-.. _mlx-examples: https://github.com/ml-explore/mlx-examples/tree/main/llms/llama
@@ -32,7 +32,7 @@ are the CPU and GPU.
   install

 .. toctree::
-   :caption: Usage
+   :caption: Usage 
   :maxdepth: 1

   usage/quick_start
@@ -45,7 +45,6 @@ are the CPU and GPU.
   usage/numpy
   usage/distributed
   usage/using_streams
-   usage/export

 .. toctree::
   :caption: Examples
@@ -54,8 +53,6 @@ are the CPU and GPU.
   examples/linear_regression
   examples/mlp
   examples/llama-inference
-   examples/data_parallelism
-   examples/tensor_parallelism

 .. toctree::
   :caption: Python API Reference
@@ -64,7 +61,6 @@ are the CPU and GPU.
   python/array
   python/data_types
   python/devices_and_streams
-   python/export
   python/ops
   python/random
   python/transforms
@@ -72,13 +68,10 @@ are the CPU and GPU.
   python/fft
   python/linalg
   python/metal
-   python/cuda
-   python/memory_management
   python/nn
   python/optimizers
   python/distributed
   python/tree_utils
-   python/printoptions

 .. toctree::
   :caption: C++ API Reference
@@ -92,6 +85,4 @@ are the CPU and GPU.

   dev/extensions
   dev/metal_debugger
-   dev/metal_logging
   dev/custom_metal_kernels
-   dev/mlx_in_cpp
@@ -1,5 +1,3 @@
-.. _build_and_install:
-
 Build and Install
 =================

@@ -13,51 +11,22 @@ silicon computer is

    pip install mlx

-To install from PyPI your system must meet the following requirements:
+To install from PyPI you must meet the following requirements:

- Using `Apple silicon <https://support.apple.com/en-us/116943>`_
- Using a native Python >= 3.10
- macOS >= 14.0
+- Using an M series chip (Apple silicon)
+- Using a native Python >= 3.9
+- macOS >= 13.5

 .. note::
-    MLX is only available on devices running macOS >= 14.0 and higher.
+    MLX is only available on devices running macOS >= 13.5
+    It is highly recommended to use macOS 14 (Sonoma)

-CUDA
-^^^^

-MLX has a CUDA backend which you can install with:
+MLX is also available on conda-forge. To install MLX with conda do:

 .. code-block:: shell

-    pip install mlx[cuda12]
-
-
-To install the CUDA package from PyPi your system must meet the following
-requirements:
-
- Nvidia architecture >= SM 7.5
- Nvidia driver >= 550.54.14
- CUDA toolkit >= 12.0
- Linux distribution with glibc >= 2.35
- Python >= 3.10
-
-For CUDA 13 use ``pip install mlx[cuda13]``. The CUDA 13 package requires
-an Nvidia driver >= 580 or an appropriate CUDA compatibility package.
-
-CPU-only (Linux)
-^^^^^^^^^^^^^^^^
-
-For a CPU-only version of MLX that runs on Linux use:
-
-.. code-block:: shell
-
-    pip install mlx[cpu]
-
-To install the CPU-only package from PyPi your system must meet the following
-requirements:
-
- Linux distribution with glibc >= 2.35
- Python >= 3.10
+   conda install conda-forge::mlx


 Troubleshooting
@@ -83,9 +52,8 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^

- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux)
- A C++ compiler with C++20 support (e.g. Clang >= 15.0)
- `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
+- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
+- `cmake <https://cmake.org/>`_ -- version 3.24 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0

 .. note::
@@ -95,8 +63,6 @@ Build Requirements
 Python API
 ^^^^^^^^^^

-.. _python install:
-
 To build and install the MLX python library from source, first, clone MLX from
 `its GitHub repo <https://github.com/ml-explore/mlx>`_:

@@ -108,20 +74,20 @@ Then simply build and install MLX using pip:

 .. code-block:: shell

-  pip install .
+  CMAKE_BUILD_PARALLEL_LEVEL=8 pip install .

 For developing, install the package with development dependencies, and use an
 editable install:

 .. code-block:: shell

-  pip install -e ".[dev]"
+  CMAKE_BUILD_PARALLEL_LEVEL=8 pip install -e ".[dev]"

 Once the development dependencies are installed, you can build faster with:

 .. code-block:: shell

- python setup.py build_ext --inplace
+ CMAKE_BUILD_PARALLEL_LEVEL=8 python setup.py build_ext --inplace

 Run the tests with:

@@ -129,11 +95,16 @@ Run the tests with:

  python -m unittest discover python/tests

+Optional: Install stubs to enable auto completions and type checking from your
+IDE:
+
+.. code-block:: shell
+
+  python setup.py generate_stubs
+
 C++ API
 ^^^^^^^

-.. _cpp install:
-
 Currently, MLX must be built and installed from source.

 Similarly to the python library, to build and install the MLX C++ library start
@@ -212,7 +183,6 @@ should point to the path to the built metal library.

      xcrun -sdk macosx --show-sdk-version

-
 Binary Size Minimization
 ~~~~~~~~~~~~~~~~~~~~~~~~

@@ -241,50 +211,6 @@ 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 across reboots.

-Linux
-^^^^^
-
-To build from source on Linux (CPU only), install the BLAS and LAPACK headers.
-For example on Ubuntu, run the following:
-
-.. code-block:: shell
-
-   apt-get update -y
-   apt-get install libblas-dev liblapack-dev liblapacke-dev -y
-
-From here follow the instructions to install either the :ref:`Python <python
-install>` or :ref:`C++ <cpp install>` APIs.
-
-CUDA
-^^^^
-
-To build from source on Linux with CUDA, install the BLAS and LAPACK headers
-and the CUDA toolkit. For example on Ubuntu, run the following:
-
-.. code-block:: shell
-
-   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
-   dpkg -i cuda-keyring_1.1-1_all.deb
-   apt-get update -y
-   apt-get -y install cuda-toolkit-12-9
-   apt-get install libblas-dev liblapack-dev liblapacke-dev libcudnn9-dev-cuda-12 -y
-
-
-When building either the Python or C++ APIs make sure to pass the cmake flag
-``MLX_BUILD_CUDA=ON``. For example, to build the Python API run:
-
-.. code-block:: shell
-
-  CMAKE_ARGS="-DMLX_BUILD_CUDA=ON" pip install -e ".[dev]"
-
-To build the C++ package run:
-
-.. code-block:: shell
-
-   mkdir -p build && cd build
-   cmake .. -DMLX_BUILD_CUDA=ON && make -j
-
-
 Troubleshooting
 ^^^^^^^^^^^^^^^

@@ -19,8 +19,6 @@ Array
    array.ndim
    array.shape
    array.size
-    array.real
-    array.imag
    array.abs
    array.all
    array.any
@@ -40,7 +38,6 @@ Array
    array.log10
    array.log1p
    array.log2
-    array.logcumsumexp
    array.logsumexp
    array.max
    array.mean
@@ -1,9 +0,0 @@
-CUDA
-=====
-
-.. currentmodule:: mlx.core.cuda
-
-.. autosummary::
-  :toctree: _autosummary
-
-  is_available
@@ -51,20 +51,11 @@ The default floating point type is ``float32`` and the default integer type is
   * - ``float32``
     - 4 
     - 32-bit float
-   * - ``float64``
-     - 8
-     - 64-bit double
   * - ``complex64``
     - 8 
     - 64-bit complex float


-.. note::
-
-    Arrays with type ``float64`` only work with CPU operations. Using
-    ``float64`` arrays on the GPU will result in an exception.
-
-
 Data type are aranged in a hierarchy. See the :obj:`DtypeCategory` object
 documentation for more information. Use :func:`issubdtype` to determine if one
 ``dtype`` (or category) is a subtype of another category.
@@ -75,4 +66,3 @@ documentation for more information. Use :func:`issubdtype` to determine if one
   Dtype
   DtypeCategory
   issubdtype
-   finfo
@@ -14,10 +14,6 @@ Devices and Streams
   set_default_device
   default_stream
   new_stream
-   new_thread_local_stream
   set_default_stream
   stream
   synchronize
-   clear_streams
-   device_count
-   device_info
@@ -1,14 +0,0 @@
-.. _export:
-
-Export Functions
-================
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-   export_function
-   import_function
-   exporter
-   export_to_dot
@@ -13,4 +13,3 @@ Fast
  rope
  scaled_dot_product_attention
  metal_kernel
-  cuda_kernel
@@ -20,7 +20,3 @@ FFT
  irfft2
  rfftn
  irfftn
-  fftfreq
-  rfftfreq
-  fftshift
-  ifftshift
@@ -5,8 +5,8 @@ Linear Algebra

 .. currentmodule:: mlx.core.linalg

-.. autosummary::
-   :toctree: _autosummary
+.. autosummary:: 
+   :toctree: _autosummary 

    inv
    tri_inv
@@ -14,16 +14,7 @@ Linear Algebra
    cholesky
    cholesky_inv
    cross
-    det
    qr
    svd
-    eigvals
-    eig
    eigvalsh
    eigh
-    lu
-    lu_factor
-    pinv
-    slogdet
-    solve
-    solve_triangular
@@ -1,16 +0,0 @@
-Memory Management
-=================
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-  get_active_memory
-  get_peak_memory
-  reset_peak_memory
-  get_cache_memory
-  set_memory_limit
-  set_cache_limit
-  set_wired_limit
-  clear_cache
@@ -8,5 +8,13 @@ Metal

  is_available
  device_info
+  get_active_memory
+  get_peak_memory
+  reset_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit
+  set_wired_limit
+  clear_cache
  start_capture
  stop_capture
@@ -174,8 +174,6 @@ In detail:

   value_and_grad
   quantize
-   average_gradients
-   fsdp_apply_gradients

 .. toctree::

@@ -184,4 +182,3 @@ In detail:
   nn/functions
   nn/losses
   nn/init
-   nn/distributed
@@ -1,30 +0,0 @@
-.. _nn_distributed:
-
-Distributed
-----------
-
-Helper Routines
-^^^^^^^^^^^^^^^
-
-The :code:`mlx.nn.layers.distributed` package contains helpful routines to 
-create sharded layers from existing :class:`Modules <mlx.nn.Module>`.
-
-.. currentmodule:: mlx.nn.layers.distributed
-.. autosummary::
-   :toctree: _autosummary
-
-   shard_linear
-   shard_inplace
-
-Layers
-^^^^^^
-
-.. currentmodule:: mlx.nn
-.. autosummary::
-   :toctree: _autosummary
-   :template: nn-module-template.rst
-
-   AllToShardedLinear
-   ShardedToAllLinear
-   QuantizedAllToShardedLinear
-   QuantizedShardedToAllLinear
@@ -27,7 +27,6 @@ simple functions.
   mish
   prelu
   relu
-   relu2
   relu6
   selu
   sigmoid
@@ -10,7 +10,6 @@ Layers
   :template: nn-module-template.rst

   ALiBi
-   AllToShardedLinear
   AvgPool1d
   AvgPool2d
   AvgPool3d
@@ -47,19 +46,15 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
-   QuantizedAllToShardedLinear
   QuantizedEmbedding
   QuantizedLinear
-   QuantizedShardedToAllLinear
   RMSNorm
   ReLU
-   ReLU2
   ReLU6
   RNN
   RoPE
   SELU
   Sequential
-   ShardedToAllLinear
   Sigmoid
   SiLU
   SinusoidalPositionalEncoding
@@ -32,16 +32,13 @@ Operations
   atleast_2d
   atleast_3d
   bitwise_and
-   bitwise_invert
   bitwise_or
   bitwise_xor
   block_masked_mm
-   broadcast_arrays
   broadcast_to
   ceil
   clip
   concatenate
-   contiguous
   conj
   conjugate
   convolve
@@ -92,7 +89,6 @@ Operations
   isneginf
   isposinf
   issubdtype
-   kron
   left_shift
   less
   less_equal
@@ -103,7 +99,6 @@ Operations
   log10
   log1p
   logaddexp
-   logcumsumexp
   logical_not
   logical_and
   logical_or
@@ -112,7 +107,6 @@ Operations
   max
   maximum
   mean
-   median
   meshgrid
   min
   minimum
@@ -150,8 +144,6 @@ Operations
   sign
   sin
   sinh
-   slice
-   slice_update
   softmax
   sort
   split
@@ -176,7 +168,6 @@ Operations
   tri
   tril
   triu
-   unflatten
   var
   view
   where
@@ -51,14 +51,14 @@ the saved state. Here's a simple example:
   optimizer.update(model, grads)

   # Save the state
-   state = tree_flatten(optimizer.state, destination={})
-   mx.save_safetensors("optimizer.safetensors", state)
+   state = tree_flatten(optimizer.state)
+   mx.save_safetensors("optimizer.safetensors", dict(state))

   # Later on, for example when loading from a checkpoint,
   # recreate the optimizer and load the state
   optimizer = optim.Adam(learning_rate=1e-2)

-   state = tree_unflatten(mx.load("optimizer.safetensors"))
+   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
   optimizer.state = state

 Note, not every optimizer configuation parameter is saved in the state. For
@@ -18,5 +18,3 @@ Common Optimizers
   AdamW
   Adamax
   Lion
-   MultiOptimizer
-   Muon
@@ -1,12 +0,0 @@
-Print Options
-===============
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-  PrintOptions
-  set_printoptions
-  printoptions
-  get_printoptions
@@ -9,9 +9,7 @@ Transforms
  :toctree: _autosummary

   eval
-   async_eval
   compile
-   checkpoint
   custom_function
   disable_compile
   enable_compile
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Alex Barron	82a956c1d9	fix test	2024-12-06 10:26:54 -08:00
Alex Barron	769704653a	cpu fallback	2024-12-06 01:22:50 -08:00
Alex Barron	c89ddf62b4	add checks	2024-12-06 01:09:00 -08:00
Alex Barron	3507c104a5	add test	2024-12-06 00:45:01 -08:00
Alex Barron	12a4d89a7c	working qsdpa	2024-12-06 00:21:05 -08:00