Use uv in macOS CI

[CUDA] Fix half type matmul in cutlass kernels (#3469 )
Clearer error when shape dimension overflows int32 (#3425 )
2026-05-06 15:41:43 +09:00 · 2026-05-06 08:35:53 +09:00 · 2026-05-05 09:53:36 +09:00 · 2026-05-05 09:22:23 +09:00 · 2026-05-05 09:06:23 +09:00 · 2026-04-28 09:39:56 -07:00
1754 changed files with 47567 additions and 597559 deletions
@@ -1,579 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
-
-parameters:
-  nightly_build:
-    type: boolean
-    default: false
-  test_release:
-    type: boolean
-    default: false
-
-jobs:
-  build_documentation:
-    parameters:
-      upload-docs:
-        type: boolean
-        default: false
-    macos:
-      xcode: "26.0.0"
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            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
-            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:
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    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: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            python -m unittest discover python/tests -v
-            mpirun --bind-to none -host localhost:8 -np 8 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
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            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: "26.0.0"
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
-              brew install openmpi uv
-      - run:
-          name: Install Python package
-          command: |
-            uv venv --python 3.9
-            uv pip install \
-              nanobind==2.4.0 \
-              cmake \
-              numpy \
-              torch \
-              tensorflow \
-              unittest-xml-reporting
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/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
-            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
-      - run:
-          name: Build example extension
-          command: |
-            source .venv/bin/activate
-            cd examples/extensions
-            uv pip install -r requirements.txt
-            uv run --no-project setup.py build_ext --inplace
-            uv run --no-project python test.py
-      - store_test_results:
-          path: test-results
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/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 .venv/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: |
-            CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-              uv pip install -e . -v
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
-              METAL_DEBUG_ERROR_MODE=0 \
-              uv run --no-project python -m xmlrunner discover \
-                -v python/tests \
-                -o test-results/gpu_jit
-
-  cuda_build_and_test:
-    parameters:
-      image_date:
-        type: string
-        default: "2023.11.1"
-    machine:
-      image: "linux-cuda-12:<< parameters.image_date >>"
-      resource_class: gpu.nvidia.small.gen2
-    steps:
-      - checkout
-      - restore_cache:
-          keys:
-            - cuda-<< parameters.image_date >>-{{ arch }}-
-      - run:
-          name: Install dependencies
-          command: |
-            sudo apt-get update
-            sudo apt-get install libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install libnccl2 libnccl-dev
-            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
-            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
-            rm -rf ccache-4.11.3-linux-x86_64
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Set CCache size
-          command: ccache --max-size 1G
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
-            LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            cmake . -B build \
-              -DMLX_BUILD_CUDA=ON \
-              -DCMAKE_CUDA_COMPILER=`which nvcc` \
-              -DCMAKE_BUILD_TYPE=DEBUG
-            cmake --build build -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-      - run:
-          name: CCache report
-          command: |
-            ccache --show-stats
-            ccache --zero-stats
-            ccache --cleanup
-      - save_cache:
-          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
-          paths:
-            - /home/circleci/.cache/ccache
-
-  build_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.9"
-      xcode_version:
-        type: string
-        default: "26.0.0"
-      build_env:
-        type: string
-        default: ""
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    resource_class: m4pro.medium
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            mkdir -p ~/miniconda3
-            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o ~/miniconda3/miniconda.sh
-            bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
-            rm ~/miniconda3/miniconda.sh
-            source ~/miniconda3/bin/activate
-            conda init --all
-            conda create -n env python=<< parameters.python_version >> -y
-            conda activate env
-            pip install --upgrade cmake
-            pip install nanobind==2.4.0
-            pip install --upgrade setuptools
-            pip install numpy
-            pip install twine
-            pip install build
-      - run:
-          name: Install Python package
-          command: |
-            conda activate env
-            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
-              pip install . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            conda activate env
-            pip install typing_extensions
-            python setup.py generate_stubs
-      - run:
-          name: Build Python package
-          command: |
-            conda activate env
-            python setup.py clean --all
-            << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
-      - when:
-          condition:
-            equal: ["3.9", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  conda activate env
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 python -m build -w
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  conda activate env
-                  twine upload dist/*
-      - store_artifacts:
-          path: dist/
-
-  build_linux_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.9"
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            PYTHON=python<< parameters.python_version >>
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            TZ=Etc/UTC sudo apt-get -y install tzdata
-            sudo add-apt-repository -y ppa:deadsnakes/ppa
-            sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            $PYTHON -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            << parameters.build_env >> pip install ".[dev]" -v
-            pip install typing_extensions
-            python setup.py generate_stubs
-            python setup.py clean --all
-            MLX_BUILD_STAGE=1 << parameters.build_env >> python -m build -w
-            bash python/scripts/repair_linux.sh
-      - when:
-          condition:
-            equal: ["3.9", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  source env/bin/activate
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 \
-                    python -m build -w
-                  auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload packages
-                command: |
-                  source env/bin/activate
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-  build_cuda_release:
-    parameters:
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: xlarge
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/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 cuda-toolkit-12-9 libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install zip
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
-            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-            << parameters.build_env >> MLX_BUILD_STAGE=2 \
-              CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              python -m build -w
-            bash python/scripts/repair_cuda.sh
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  twine upload wheelhouse/*.whl
-      - 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.test_release >>
-    jobs:
-      - mac_build_and_test:
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test
-      - cuda_build_and_test:
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-      - build_documentation 
-
-  build_pypi_release:
-    when:
-      and:
-        - not: << pipeline.parameters.nightly_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"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["PYPI_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_documentation:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          upload-docs: true
-      - build_linux_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              build_env: ["PYPI_RELEASE=1"]
-      - build_cuda_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              build_env: ["PYPI_RELEASE=1"]
-
-  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:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test:
-          requires: [ hold ]
-      - cuda_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-  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"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-      - build_cuda_release
-
-  build_dev_release:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["DEV_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-              build_env: ["DEV_RELEASE=1"]
-      - build_cuda_release:
-          matrix:
-            parameters:
-              build_env: ["DEV_RELEASE=1"]
@@ -0,0 +1,31 @@
+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*
@@ -0,0 +1,38 @@
+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
@@ -0,0 +1,42 @@
+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 }}
@@ -0,0 +1,38 @@
+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)
@@ -0,0 +1,36 @@
+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: |
+        uv 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
@@ -0,0 +1,96 @@
+name: 'Build and Test on macOS'
+description: 'Build and test MLX on macOS'
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install Python package
+      env:
+        DEBUG: 1
+        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+      shell: bash
+      run: |
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
+
+    - name: Install tests dependencies
+      shell: bash
+      run: |
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
+
+    - name: Run Python tests
+      shell: bash
+      env:
+        LOW_MEMORY: 1
+      run: |
+        echo "::group::Run Python tests"
+        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
+        echo "::endgroup::"
+
+    - name: Build example extension
+      shell: bash
+      run: |
+        echo "::group::Build example extension"
+        cd examples/extensions
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
+    - name: Build CPP only
+      shell: bash
+      run: |
+        echo "::group::Build CPP only"
+        mkdir -p build
+        cd build
+        cmake ..
+        make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
+
+    - name: Run CPP tests
+      shell: bash
+      env:
+        DEVICE: gpu
+        METAL_DEVICE_WRAPPER_TYPE: 1
+        METAL_DEBUG_ERROR_MODE: 0
+      run: |
+        echo "::group::Run CPP tests"
+        ./build/tests/tests
+        ./build/tests/test_teardown
+        echo "::endgroup::"
+
+    - name: Build small binary with JIT
+      shell: bash
+      run: |
+        echo "::group::Build small binary with JIT"
+        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)
+        echo "::endgroup::"
+
+    - name: Run Python tests with JIT
+      shell: bash
+      env:
+        LOW_MEMORY: 1
+        DEVICE: gpu
+        METAL_DEVICE_WRAPPER_TYPE: 1
+        METAL_DEBUG_ERROR_MODE: 0
+      run: |
+        echo "::group::Run Python tests with JIT"
+        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
+          uv pip install -e . -v
+        python -m unittest discover -v python/tests
+        echo "::endgroup::"
@@ -0,0 +1,26 @@
+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%
@@ -0,0 +1,102 @@
+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::"
@@ -0,0 +1,32 @@
+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: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        uv venv --python ${{ inputs.python-version }}
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"
@@ -0,0 +1,42 @@
+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
@@ -0,0 +1,69 @@
+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::"
@@ -0,0 +1,21 @@
+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::"
@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
@@ -0,0 +1,48 @@
+#!/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
@@ -0,0 +1,27 @@
+#!/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
@@ -0,0 +1,152 @@
+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
@@ -0,0 +1,28 @@
+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
@@ -0,0 +1,108 @@
+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
@@ -1,20 +0,0 @@
-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
@@ -0,0 +1,251 @@
+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 Python package
+        run: uv 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,16 +3,12 @@ __pycache__/
 *.py[cod]
 *$py.class

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

 # Metal libraries
 *.metallib
-venv/

 # Distribution / packaging
 python/mlx/core
@@ -30,6 +26,7 @@ lib64/
 parts/
 sdist/
 var/
+venv/
 wheels/
 share/python-wheels/
 *.egg-info/
@@ -37,12 +34,7 @@ share/python-wheels/
 *.egg
 MANIFEST
 uv.lock
-
-# vim
-*.swp
-
-# Ignore build dir
-build/
+.DS_Store

 # Prerequisites
 *.d
@@ -52,6 +44,7 @@ build/
 *.lo
 *.o
 *.obj
+*.ilk

 # Precompiled Headers
 *.gch
@@ -80,9 +73,9 @@ build/
 # Debug symbols
 *.pdb

-# VSCode 
+# VSCode
 .vscode/
-.DS_Store
-
 # Jetbrains
-.cache
+.cache/
+# vim
+*.swp
@@ -1,16 +1,22 @@
 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: v19.1.7
+    rev: v21.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: 25.1.0
+    rev: 26.1.0
    hooks:
    -   id: black
    
 -   repo: https://github.com/pycqa/isort
-    rev: 6.0.0
+    rev: 7.0.0
    hooks:
    -   id: isort
        args:
@@ -22,7 +22,7 @@ project(

 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
@@ -40,11 +40,14 @@ option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
-option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
+option(MLX_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)

 # --------------------- Processor tests -------------------------
 message(
@@ -74,12 +77,70 @@ 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()
+endif()
+
 # ----------------------------- Lib -----------------------------

 include(FetchContent)
@@ -88,8 +149,17 @@ 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()
 endif()

 if(MLX_BUILD_METAL)
@@ -122,9 +192,12 @@ 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_macOS15_iOS18.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_26.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(
@@ -133,7 +206,6 @@ 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}>
@@ -151,14 +223,17 @@ if(WIN32)
  if(MSVC)
    # GGUF does not build with MSVC.
    set(MLX_BUILD_GGUF OFF)
-    # There is no prebuilt OpenBLAS distribution for MSVC.
-    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
+  endif()
+  # 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.1
+    GIT_TAG v1.4.2
    EXCLUDE_FROM_ALL)
  block()
  set(BUILD_SHARED_LIBS OFF)
@@ -182,20 +257,25 @@ if(MLX_BUILD_CPU)
    target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
    add_compile_definitions(MLX_USE_ACCELERATE)
    add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
-    # Download and build OpenBLAS from source code.
+  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
-      GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git
-      GIT_TAG v0.3.28
-      EXCLUDE_FROM_ALL)
-    set(BUILD_STATIC_LIBS ON) # link statically
-    set(NOFORTRAN ON) # msvc has no fortran compiler
+      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)
-    target_include_directories(
-      mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include"
-                  "${CMAKE_BINARY_DIR}/generated" "${CMAKE_BINARY_DIR}")
+    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()
    if(${CMAKE_HOST_APPLE})
      # The blas shipped in macOS SDK is not supported, search homebrew for
@@ -241,22 +321,28 @@ 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)
+endif()
+
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)

 target_include_directories(
  mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
             $<INSTALL_INTERFACE:include>)

-# Do not add mlx_EXPORTS define for shared library.
-set_target_properties(mlx PROPERTIES DEFINE_SYMBOL "")
-
 if(USE_SYSTEM_FMT)
  find_package(fmt REQUIRED)
 else()
  FetchContent_Declare(
    fmt
    GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-    GIT_TAG 10.2.1
+    GIT_TAG 12.1.0
    EXCLUDE_FROM_ALL)
  FetchContent_MakeAvailable(fmt)
 endif()
@@ -265,14 +351,16 @@ 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.8
+    Python 3.10
    COMPONENTS Interpreter Development.Module
    REQUIRED)
-  execute_process(
-    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE nanobind_ROOT)
-  find_package(nanobind CONFIG 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)
  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()

@@ -292,6 +380,15 @@ 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
@@ -75,7 +75,7 @@ void time_irregular_binary_ops_3D() {

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

@@ -115,7 +115,7 @@ void time_irregular_binary_ops_4D() {

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

  auto a = mx::random::uniform(shape);
  TIMEM("2D regular", mx::astype, a, mx::int32, device);
@@ -1,6 +1,5 @@
 # Copyright © 2023 Apple Inc.

-import argparse
 import os
 import subprocess
 import time
@@ -0,0 +1,193 @@
+# 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.time()
+    s = time.perf_counter()
    for i in range(100):
        f(*args)
-    e = time.time()
+    e = time.perf_counter()
    return e - s


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

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


@@ -0,0 +1,152 @@
+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}"
+            )
@@ -0,0 +1,119 @@
+# 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"
+            )
@@ -0,0 +1,236 @@
+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()
@@ -176,6 +176,8 @@ if __name__ == "__main__":
          (  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),
    )

    shapes_80 = (
@@ -183,6 +185,8 @@ if __name__ == "__main__":
          (  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),
    )

    shapes_128 = (
@@ -190,6 +194,8 @@ if __name__ == "__main__":
          (  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),
    )
    # fmt: on

@@ -0,0 +1,209 @@
+# 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()
@@ -0,0 +1,109 @@
+# 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}"
+            )
@@ -31,8 +31,8 @@ def measure_runtime(fn, **kwargs):
    for _ in range(5):
        fn(**kwargs)

-    tic = time.time()
+    tic = time.perf_counter()
    iters = 100
    for _ in range(iters):
        fn(**kwargs)
-    return (time.time() - tic) * 1000 / iters
+    return (time.perf_counter() - tic) * 1000 / iters
@@ -0,0 +1,177 @@
+# 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()
@@ -0,0 +1,3 @@
+# 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.
@@ -26,6 +26,7 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
+PREDEFINED             = MLX_API=

 ################################################################################
 # Compound extraction control.                                                 #
@@ -38,3 +38,17 @@ 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,4 +0,0 @@
-# Sphinx build info version 1
-# This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: a85b8c51cc3365be6cea33d5b2a3c4c7
-tags: 645f666f9bcd5a90fca523b33c5a78b7
@@ -1,7 +0,0 @@
-.. _cpp_ops:
-
-Operations
-==========
-
-.. doxygengroup:: ops
-   :content-only:
@@ -1,445 +0,0 @@
-.. _custom_metal_kernels:
-
-Custom Metal Kernels
-====================
-
-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):
-      outputs = kernel(
-          inputs=[a],
-          template=[("T", mx.float32)],
-          grid=(a.size, 1, 1),
-          threadgroup=(256, 1, 1),
-          output_shapes=[a.shape],
-          output_dtypes=[a.dtype],
-      )
-      return outputs[0]
-
-  a = mx.random.normal(shape=(4, 16)).astype(mx.float16)
-  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.
-
-The full function signature will be generated using:
-
-* The shapes/dtypes of ``inputs``
-    In the above, ``a`` is an ``mx.array`` of type ``mx.float16`` and we pass it with the key ``inp``
-    so we will add ``const device float16_t* inp`` to the signature.
-    ``inp_shape``, ``inp_strides`` and ``inp_ndim`` are also added for convenience if they are present
-    in ``source``.
-* The list of ``output_dtypes``
-    In the above, ``out`` is an ``mx.array`` of type ``mx.float16``
-    so we add ``device float16_t* out``.
-* Template parameters passed using ``template``
-    In the above, ``template=[("T", mx.float32)]`` adds a template of ``template <typename T>`` to the function
-    and instantiates the template with ``custom_kernel_myexp_float<float>``.
-    Template parameters can be ``mx.core.Dtype``, ``int`` or ``bool``.
-* Metal attributes used in ``source`` such as ``[[thread_position_in_grid]]``
-    These will be added as function arguments.
-    All the attributes defined in Table 5.8 of the `Metal Shading Language Specification <https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf>`_ are supported.
-
-Putting this all together, the generated function signature for ``myexp`` is as follows:
-
-.. code-block:: cpp
-
-  template <typename T>
-  [[kernel]] void custom_kernel_myexp_float(
-    const device float16_t* inp [[buffer(0)]],
-    device float16_t* out [[buffer(1)]],
-    uint3 thread_position_in_grid [[thread_position_in_grid]]) {
-
-          uint elem = thread_position_in_grid.x;
-          T tmp = inp[elem];
-          out[elem] = metal::exp(tmp);
-
-  }
-
-  template [[host_name("custom_kernel_myexp_float")]] [[kernel]] decltype(custom_kernel_myexp_float<float>) custom_kernel_myexp_float<float>;
-
-Note: ``grid`` and ``threadgroup`` are parameters to the Metal `dispatchThreads
-<https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads>`_
-function. This means we will launch ``mx.prod(grid)`` threads, subdivided into
-``threadgroup`` size threadgroups.  For optimal performance, each thread group
-dimension should be less than or equal to the corresponding grid dimension.
-
-Passing ``verbose=True`` to :func:`ast.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.
-
-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.
-
-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):
-      outputs = kernel(
-          inputs=[a],
-          template=[("T", mx.float32)],
-          grid=(a.size, 1, 1),
-          threadgroup=(256, 1, 1),
-          output_shapes=[a.shape],
-          output_dtypes=[a.dtype],
-      )
-      return outputs[0]
-
-  a = mx.random.normal(shape=(4, 16)).astype(mx.float16)
-  # make non-contiguous
-  a = a[::2]
-  b = exp_elementwise(a)
-  assert mx.allclose(b, mx.exp(a))
-
-Complex Example
-----------------------------
-
-Let's implement a more complex example: ``grid_sample`` in ``"bilinear"`` mode.
-
-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
-
-      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_sw = ix_nw
-      iy_sw = 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)
-
-      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)
-
-      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
-
-      return output
-
-Now let's use :func:`custom_function` together with :func:`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];
-
-      int w_stride = C;
-      int h_stride = W * w_stride;
-      int b_stride = H * h_stride;
-
-      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;
-
-      int ix_nw = floor(ix);
-      int iy_nw = floor(iy);
-
-      int ix_ne = ix_nw + 1;
-      int iy_ne = iy_nw;
-
-      int ix_sw = ix_nw;
-      int iy_sw = iy_nw + 1;
-
-      int ix_se = ix_nw + 1;
-      int iy_se = iy_nw + 1;
-
-      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 batch_idx = elem / C / gH / gW * b_stride;
-      int channel_idx = elem % C;
-      int base_idx = batch_idx + channel_idx;
-
-      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];
-
-      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;
-
-      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,
-  )
-
-  @mx.custom_function
-  def grid_sample(x, grid):
-
-      assert x.ndim == 4, "`x` must be 4D."
-      assert grid.ndim == 4, "`grid` must be 4D."
-
-      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]
-
-For a reasonably sized input such as:
-
-.. code-block:: python
-
-  x.shape = (8, 1024, 1024, 64)
-  grid.shape = (8, 256, 256, 2)
-
-On an M1 Max, we see a big performance improvement:
-
-``55.7ms -> 6.7ms => 8x speed up``
-
-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.
-
-The backwards pass requires atomically updating ``x_grad``/``grid_grad`` and so
-requires a few extra :func:`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.
-
-* ``atomic_outputs=True``
-    Designate all of the kernel outputs as ``atomic`` in the function signature. 
-    This means we can use Metal's ``atomic`` features to simultaneously update the ``x_grad`` and ``grid_grad`` arrays from multiple threadgroups. 
-    See section 6.15 of the `Metal Shading Language Specification <https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf>`_ for more details.
-
-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;
-
-      int gH = grid_shape[1];
-      int gW = grid_shape[2];
-
-      int w_stride = C;
-      int h_stride = W * w_stride;
-      int b_stride = H * h_stride;
-
-      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_nw = floor(ix);
-      int iy_nw = floor(iy);
-
-      int ix_ne = ix_nw + 1;
-      int iy_ne = iy_nw;
-
-      int ix_sw = ix_nw;
-      int iy_sw = iy_nw + 1;
-
-      int ix_se = ix_nw + 1;
-      int iy_se = iy_nw + 1;
-
-      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 batch_idx = elem / C_padded / gH / gW * b_stride;
-      int channel_idx = elem % C_padded;
-      int base_idx = batch_idx + channel_idx;
-
-      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_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_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 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_se = x[offset];
-              gix += I_se * (iy - iy_nw) * cot;
-              giy += I_se * (ix - ix_nw) * cot;
-          }
-      }
-
-      T gix_mult = W / 2;
-      T giy_mult = H / 2;
-
-      // Reduce across each simdgroup first.
-      // This is much faster than relying purely on atomics.
-      gix = simd_sum(gix);
-      giy = simd_sum(giy);
-
-      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,
-  )
-
-  @grid_sample.vjp
-  def grid_sample_vjp(primals, cotangent, _):
-      x, grid = primals
-      B, _, _, C = x.shape
-      _, gN, gM, D = grid.shape
-
-      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]
-
-There's an even larger speed up for the vjp:
-
-``676.4ms -> 16.7ms => 40x speed up``
@@ -1,811 +0,0 @@
-Custom Extensions in MLX
-========================
-
-You can extend MLX with custom operations on the CPU or GPU. This guide
-explains how to do that with a simple example.
-
-Introducing the Example
-----------------------
-
-Let's say you would like an operation that takes in two arrays, ``x`` and
-``y``, scales them both by coefficients ``alpha`` and ``beta`` respectively,
-and then adds them together to get the result ``z = alpha * x + beta * y``.
-You can do that in MLX directly:
-
-.. code-block:: python
-
-    import mlx.core as mx
-
-    def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
-        return alpha * x + beta * y
-
-This function performs that operation while leaving the implementation and
-function transformations to MLX.
-
-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:
-
-* The structure of the MLX library.
-* Implementing a CPU operation.
-* Implementing a GPU operation using metal.
-* Adding the ``vjp`` and ``jvp`` function transformation.
-* Building a custom extension and binding it to python.
-
-Operations and Primitives
-------------------------
-
-Operations in MLX build the computation graph. Primitives provide the rules for
-evaluating and transforming the graph. Let's start by discussing operations in
-more detail.
-
-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
-``y``, and two scalars, ``alpha`` and ``beta``. This is how to define it in
-C++:
-
-.. code-block:: C++
-
-    /**
-    *  Scale and sum two vectors element-wise
-    *  z = alpha * x + beta * y
-    *
-    *  Use NumPy-style broadcasting between x and y
-    *  Inputs are upcasted to floats if needed
-    **/
-    array axpby(
-        const array& x, // Input array x
-        const array& y, // Input array y
-        const float alpha, // Scaling factor for x
-        const float beta, // Scaling factor for y
-        StreamOrDevice s = {} // Stream on which to schedule the operation
-    );
-
-The simplest way to implement this is with existing operations:
-
-.. code-block:: C++
-
-    array axpby(
-        const array& x, // Input array x
-        const array& y, // Input array y
-        const float alpha, // Scaling factor for x
-        const float beta, // Scaling factor for y
-        StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
-    ) {
-        // Scale x and y on the provided stream
-        auto ax = multiply(array(alpha), x, s);
-        auto by = multiply(array(beta), y, s);
-
-        // Add and return
-        return add(ax, by, s);
-    }
-
-The operations themselves do not contain the implementations that act on the
-data, nor do they contain the rules of transformations. Rather, they are an
-easy to use interface that use :class:`Primitive` building blocks.
-
-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
-: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
-more concrete:
-
-.. code-block:: C++
-
-    class Axpby : public Primitive {
-      public:
-        explicit Axpby(Stream stream, float alpha, float beta)
-            : Primitive(stream), alpha_(alpha), beta_(beta){};
-
-        /**
-        * A primitive must know how to evaluate itself on the CPU/GPU
-        * for the given inputs and populate the output array.
-        *
-        * To avoid unnecessary allocations, the evaluation function
-        * is responsible for allocating space for the array.
-        */
-        void eval_cpu(
-            const std::vector<array>& inputs,
-            std::vector<array>& outputs) override;
-        void eval_gpu(
-            const std::vector<array>& inputs,
-            std::vector<array>& outputs) override;
-
-        /** The Jacobian-vector product. */
-        std::vector<array> jvp(
-            const std::vector<array>& primals,
-            const std::vector<array>& tangents,
-            const std::vector<int>& argnums) override;
-
-        /** The vector-Jacobian product. */
-        std::vector<array> vjp(
-            const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
-            const std::vector<int>& argnums,
-            const std::vector<array>& outputs) override;
-
-        /**
-        * The primitive must know how to vectorize itself across
-        * the given axes. The output is a pair containing the array
-        * representing the vectorized computation and the axis which
-        * corresponds to the output vectorized dimension.
-        */
-        std::pair<std::vector<array>, std::vector<int>> vmap(
-            const std::vector<array>& inputs,
-            const std::vector<int>& axes) override;
-
-        /** The name of primitive. */
-        const char* name() const override {
-          return "Axpby";
-        }
-
-        /** Equivalence check **/
-        bool is_equivalent(const Primitive& other) const override;
-
-      private:
-        float alpha_;
-        float beta_;
-    };
-
-The :class:`Axpby` class derives from the base :class:`Primitive` class. The
-:class:`Axpby` treats ``alpha`` and ``beta`` as parameters. It then provides
-implementations of how the output array is produced given the inputs through
-:meth:`Axpby::eval_cpu` and :meth:`Axpby::eval_gpu`. It also provides rules
-of transformations in :meth:`Axpby::jvp`, :meth:`Axpby::vjp`, and
-:meth:`Axpby::vmap`.
-
-Using the Primitive
-^^^^^^^^^^^^^^^^^^^
-
-Operations can use this :class:`Primitive` to add a new :class:`array` to the
-computation graph. An :class:`array` can be constructed by providing its data
-type, shape, the :class:`Primitive` that computes it, and the :class:`array`
-inputs that are passed to the primitive.
-
-Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
-
-.. code-block:: C++
-
-    array axpby(
-        const array& x, // Input array x
-        const array& y, // Input array y
-        const float alpha, // Scaling factor for x
-        const float beta, // Scaling factor for y
-        StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
-    ) {
-        // Promote dtypes between x and y as needed
-        auto promoted_dtype = promote_types(x.dtype(), y.dtype());
-
-        // Upcast to float32 for non-floating point inputs x and y
-        auto out_dtype = issubdtype(promoted_dtype, float32)
-            ? promoted_dtype
-            : promote_types(promoted_dtype, float32);
-
-        // Cast x and y up to the determined dtype (on the same stream s)
-        auto x_casted = astype(x, out_dtype, s);
-        auto y_casted = astype(y, out_dtype, s);
-
-        // Broadcast the shapes of x and y (on the same stream s)
-        auto broadcasted_inputs = broadcast_arrays({x_casted, y_casted}, s);
-        auto out_shape = broadcasted_inputs[0].shape();
-
-        // Construct the array as the output of the Axpby primitive
-        // with the broadcasted and upcasted arrays as inputs
-        return array(
-            /* const std::vector<int>& shape = */ out_shape,
-            /* Dtype dtype = */ out_dtype,
-            /* std::unique_ptr<Primitive> primitive = */
-            std::make_shared<Axpby>(to_stream(s), alpha, beta),
-            /* const std::vector<array>& inputs = */ broadcasted_inputs);
-    }
-
-
-This operation now handles the following:
-
-#. Upcast inputs and resolve the output data type.
-#. Broadcast the inputs and resolve the output shape.
-#. Construct the primitive :class:`Axpby` using the given stream, ``alpha``, and ``beta``.
-#. Construct the output :class:`array` using the primitive and the inputs.
-
-Implementing the Primitive
--------------------------
-
-No computation happens when we call the operation alone. The operation only
-builds the computation graph. When we evaluate the output array, MLX schedules
-the execution of the computation graph, and calls :meth:`Axpby::eval_cpu` or
-:meth:`Axpby::eval_gpu` depending on the stream/device specified by the user.
-
-.. warning::
-    When :meth:`Primitive::eval_cpu` or :meth:`Primitive::eval_gpu` are called,
-    no memory has been allocated for the output array. It falls on the implementation
-    of these functions to allocate memory as needed.
-
-Implementing the CPU Back-end
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Let's start by implementing :meth:`Axpby::eval_cpu`.
-
-The 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()));
-
-    // 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);
-
-    // 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_);
-
-      // 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];
-      }
-    });
-  }
-
-Our implementation should work for all incoming floating point arrays.
-Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
-``complex64``. We throw an error if we encounter an unexpected type.
-
-.. 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];
-
-      // 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.");
-      }
-    }
-
-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.
-
-Implementing the GPU Back-end
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Apple silicon devices address their GPUs using the Metal_ shading language, and
-GPU kernels in MLX are written using Metal.
-
-.. note::
-
-    Here are some helpful resources if you are new to Metal:
-
-    * A walkthrough of the metal compute pipeline: `Metal Example`_
-    * Documentation for metal shading language: `Metal Specification`_
-    * Using metal from C++: `Metal-cpp`_
-
-Let's keep the GPU kernel simple. We will launch exactly as many threads as
-there are elements in the output. Each thread will pick the element it needs
-from ``x`` and ``y``, do the point-wise operation, and update its assigned
-element in the output.
-
-.. code-block:: C++
-
-    template <typename T>
-    [[kernel]] void axpby_general(
-            device const T* x [[buffer(0)]],
-            device const T* y [[buffer(1)]],
-            device T* out [[buffer(2)]],
-            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 int& ndim [[buffer(8)]],
-            uint index [[thread_position_in_grid]]) {
-        // Convert linear indices to offsets in array
-        auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
-        auto y_offset = elem_to_loc(index, shape, y_strides, ndim);
-
-        // Do the operation and update the output
-        out[index] =
-            static_cast<T>(alpha) * x[x_offset] + static_cast<T>(beta) * y[y_offset];
-    }
-
-We then need to instantiate this template for all floating point types and give
-each instantiation a unique host name so we can identify it.
-
-.. 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)
-
-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
-below.
-
-.. code-block:: C++
-
-    /** Evaluate primitive on GPU */
-    void Axpby::eval_gpu(
-      const std::vector<array>& inputs,
-      std::vector<array>& outputs) {
-        // Prepare inputs
-        assert(inputs.size() == 2);
-        auto& x = inputs[0];
-        auto& y = inputs[1];
-        auto& out = outputs[0];
-
-        // Each primitive carries the stream it should execute on
-        // and each stream carries its device identifiers
-        auto& s = stream();
-        // We get the needed metal device using the stream
-        auto& d = metal::device(s.device);
-
-        // Allocate output memory
-        out.set_data(allocator::malloc(out.nbytes()));
-
-        // Resolve name of kernel
-        std::stream kname;
-        kname = "axpby_general_" + type_to_name(out);
-
-        // Load the metal library
-        auto lib = d.get_library("mlx_ext", current_binary_dir());
-
-        // Make a kernel from this metal library
-        auto kernel = d.get_kernel(kname, lib);
-
-        // Prepare to encode kernel
-        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
-        // those in the kernel declaration at axpby.metal
-        int ndim = out.ndim();
-        size_t nelem = out.size();
-
-        // Encode input arrays to kernel
-        compute_encoder.set_input_array(x, 0);
-        compute_encoder.set_input_array(y, 1);
-
-        // Encode output arrays to kernel
-        compute_encoder.set_output_array(out, 2);
-
-        // Encode alpha and beta
-        compute_encoder.set_bytes(alpha_, 3);
-        compute_encoder.set_bytes(beta_, 4);
-
-        // Encode shape, strides and ndim
-        compute_encoder.set_vector_bytes(x.shape(), 5);
-        compute_encoder.set_vector_bytes(x.strides(), 6);
-        compute_encoder.set_bytes(y.strides(), 7);
-        compute_encoder.set_bytes(ndim, 8);
-
-        // We launch 1 thread for each input and make sure that the number of
-        // threads in any given threadgroup is not higher than the max allowed
-        size_t tgp_size = std::min(nelem, kernel->maxTotalThreadsPerThreadgroup());
-
-        // Fix the 3D size of each threadgroup (in terms of threads)
-        MTL::Size group_dims = MTL::Size(tgp_size, 1, 1);
-
-        // Fix the 3D size of the launch grid (in terms of threads)
-        MTL::Size grid_dims = MTL::Size(nelem, 1, 1);
-
-        // Launch the grid with the given number of threads divided among
-        // the given threadgroups
-        compute_encoder.dispatch_threads(grid_dims, group_dims);
-    }
-
-We can now call the :meth:`axpby` operation on both the CPU and the GPU!
-
-A few things to note about MLX and Metal before moving on. MLX keeps track of
-the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is
-associated. We rely on :meth:`d.get_command_encoder` to give us the active
-metal compute command encoder instead of building a new one and calling
-:meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute
-pipelines) to the active command buffer until some specified limit is hit or
-the command buffer needs to be flushed for synchronization.
-
-Primitive Transforms
-^^^^^^^^^^^^^^^^^^^^^
-
-Next, let's add implementations for transformations in a :class:`Primitive`.
-These transformations can be built on top of other operations, including the
-one we just defined:
-
-.. code-block:: C++
-
-    /** The Jacobian-vector product. */
-    std::vector<array> Axpby::jvp(
-            const std::vector<array>& primals,
-            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
-        // that are scheduled on the same stream as the primitive
-
-        // If argnums = {0}, we only push along x in which case the
-        // jvp is just the tangent scaled by alpha
-        // Similarly, if argnums = {1}, the jvp is just the tangent
-        // scaled by beta
-        if (argnums.size() > 1) {
-            auto scale = argnums[0] == 0 ? alpha_ : beta_;
-            auto scale_arr = array(scale, tangents[0].dtype());
-            return {multiply(scale_arr, tangents[0], stream())};
-        }
-        // 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())};
-        }
-    }
-
-.. code-block:: C++
-
-    /** The vector-Jacobian product. */
-    std::vector<array> Axpby::vjp(
-            const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
-            const std::vector<int>& argnums,
-            const std::vector<int>& /* unused */) {
-        // Reverse mode diff
-        std::vector<array> vjps;
-        for (auto arg : argnums) {
-            auto scale = arg == 0 ? alpha_ : beta_;
-            auto scale_arr = array(scale, cotangents[0].dtype());
-            vjps.push_back(multiply(scale_arr, cotangents[0], stream()));
-        }
-        return vjps;
-    }
-
-Note, a transformation does not need to be fully defined to start using
-the :class:`Primitive`.
-
-.. code-block:: C++
-
-    /** Vectorize primitive along given axis */
-    std::pair<std::vector<array>, std::vector<int>> Axpby::vmap(
-            const std::vector<array>& inputs,
-            const std::vector<int>& axes) {
-        throw std::runtime_error("[Axpby] vmap not implemented.");
-    }
-
-Building and Binding
--------------------
-
-Let's look at the overall directory structure first.
-
-| extensions
-| ├── axpby
-| │   ├── axpby.cpp
-| │   ├── axpby.h
-| │   └── axpby.metal
-| ├── mlx_sample_extensions
-| │   └── __init__.py
-| ├── bindings.cpp
-| ├── CMakeLists.txt
-| └── setup.py
-
-* ``extensions/axpby/`` defines the C++ extension library
-* ``extensions/mlx_sample_extensions`` sets out the structure for the
-  associated Python package
-* ``extensions/bindings.cpp`` provides Python bindings for our operation
-* ``extensions/CMakeLists.txt`` holds CMake rules to build the library and
-  Python bindings
-* ``extensions/setup.py`` holds the ``setuptools`` rules to build and install
-  the Python package
-
-Binding to Python
-^^^^^^^^^^^^^^^^^^
-
-We use nanobind_ to build a Python API for the C++ library. Since bindings for
-components such as :class:`mlx.core.array`, :class:`mlx.core.stream`, etc. are
-already provided, adding our :meth:`axpby` is simple.
-
-.. code-block:: C++
-
-   NB_MODULE(_ext, m) {
-        m.doc() = "Sample extension for MLX";
-
-        m.def(
-            "axpby",
-            &axpby,
-            "x"_a,
-            "y"_a,
-            "alpha"_a,
-            "beta"_a,
-            nb::kw_only(),
-            "stream"_a = nb::none(),
-            R"(
-                Scale and sum two vectors element-wise
-                ``z = alpha * x + beta * y``
-
-                Follows numpy style broadcasting between ``x`` and ``y``
-                Inputs are upcasted to floats if needed
-
-                Args:
-                    x (array): Input array.
-                    y (array): Input array.
-                    alpha (float): Scaling factor for ``x``.
-                    beta (float): Scaling factor for ``y``.
-
-                Returns:
-                    array: ``alpha * x + beta * y``
-            )");
-    }
-
-Most of the complexity in the above example comes from additional bells and
-whistles such as the literal names and doc-strings.
-
-.. warning::
-
-    :mod:`mlx.core` must be imported before importing
-    :mod:`mlx_sample_extensions` as defined by the nanobind module above to
-    ensure that the casters for :mod:`mlx.core` components like
-    :class:`mlx.core.array` are available.
-
-.. _Building with CMake:
-
-Building with CMake
-^^^^^^^^^^^^^^^^^^^^
-
-Building the C++ extension library only requires that you ``find_package(MLX
-CONFIG)`` and then link it to your library.
-
-.. code-block:: cmake
-
-    # Add library
-    add_library(mlx_ext)
-
-    # Add sources
-    target_sources(
-        mlx_ext
-        PUBLIC
-        ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp
-    )
-
-    # Add include headers
-    target_include_directories(
-        mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}
-    )
-
-    # Link to mlx
-    target_link_libraries(mlx_ext PUBLIC mlx)
-
-We also need to build the attached Metal library. For convenience, we provide a
-:meth:`mlx_build_metallib` function that builds a ``.metallib`` target given
-sources, headers, destinations, etc. (defined in ``cmake/extension.cmake`` and
-automatically imported with MLX package).
-
-Here is what that looks like in practice:
-
-.. code-block:: cmake
-
-    # Build metallib
-    if(MLX_BUILD_METAL)
-
-    mlx_build_metallib(
-        TARGET mlx_ext_metallib
-        TITLE mlx_ext
-        SOURCES ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
-        INCLUDE_DIRS ${PROJECT_SOURCE_DIR} ${MLX_INCLUDE_DIRS}
-        OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}
-    )
-
-    add_dependencies(
-        mlx_ext
-        mlx_ext_metallib
-    )
-
-    endif()
-
-Finally, we build the nanobind_ bindings
-
-.. code-block:: cmake
-
-    nanobind_add_module(
-      _ext
-      NB_STATIC STABLE_ABI LTO NOMINSIZE
-      NB_DOMAIN mlx
-      ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
-    )
-    target_link_libraries(_ext PRIVATE mlx_ext)
-
-    if(BUILD_SHARED_LIBS)
-      target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
-    endif()
-
-Building with ``setuptools``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Once we have set out the CMake build rules as described above, we can use the
-build utilities defined in :mod:`mlx.extension`:
-
-.. code-block:: python
-
-    from mlx import extension
-    from setuptools import setup
-
-    if __name__ == "__main__":
-        setup(
-            name="mlx_sample_extensions",
-            version="0.0.0",
-            description="Sample C++ and Metal extensions for MLX primitives.",
-            ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")],
-            cmdclass={"build_ext": extension.CMakeBuild},
-            packages=["mlx_sample_extensions"],
-            package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
-            extras_require={"dev":[]},
-            zip_safe=False,
-            python_requires=">=3.8",
-        )
-
-.. note::
-    We treat ``extensions/mlx_sample_extensions`` as the package directory
-    even though it only contains a ``__init__.py`` to ensure the following:
-
-    * :mod:`mlx.core` must be imported before importing :mod:`_ext`
-    * The C++ extension library and the metal library are co-located with the python
-      bindings and copied together if the package is installed
-
-To build the package, first install the build dependencies with ``pip install
-r requirements.txt``.  You can then build inplace for development using
-``python setup.py build_ext -j8 --inplace`` (in ``extensions/``)
-
-This results in the directory structure:
-
-| extensions
-| ├── mlx_sample_extensions
-| │   ├── __init__.py
-| │   ├── libmlx_ext.dylib # C++ extension library
-| │   ├── mlx_ext.metallib # Metal library
-| │   └── _ext.cpython-3x-darwin.so # Python Binding
-| ...
-
-When you try to install using the command ``python -m pip install .`` (in
-``extensions/``), the package will be installed with the same structure as
-``extensions/mlx_sample_extensions`` and the C++ and Metal library will be
-copied along with the Python binding since they are specified as
-``package_data``.
-
-Usage
-----
-
-After installing the extension as described above, you should be able to simply
-import the Python package and play with it as you would any other MLX operation.
-
-Let's look at a simple script and its results:
-
-.. code-block:: python
-
-    import mlx.core as mx
-    from mlx_sample_extensions import axpby
-
-    a = mx.ones((3, 4))
-    b = mx.ones((3, 4))
-    c = axpby(a, b, 4.0, 2.0, stream=mx.cpu)
-
-    print(f"c shape: {c.shape}")
-    print(f"c dtype: {c.dtype}")
-    print(f"c is correct: {mx.all(c == 6.0).item()}")
-
-Output:
-
-.. code-block::
-
-    c shape: [3, 4]
-    c dtype: float32
-    c is correct: 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.
-
-.. code-block:: python
-
-    import mlx.core as mx
-    from mlx_sample_extensions import axpby
-    import time
-
-    def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
-        return alpha * x + beta * y
-
-    M = 4096
-    N = 4096
-
-    x = mx.random.normal((M, N))
-    y = mx.random.normal((M, N))
-    alpha = 4.0
-    beta = 2.0
-
-    mx.eval(x, y)
-
-    def bench(f):
-        # Warm up
-        for i in range(5):
-            z = f(x, y, alpha, beta)
-            mx.eval(z)
-
-        # Timed run
-        s = time.time()
-        for i in range(100):
-            z = f(x, y, alpha, beta)
-            mx.eval(z)
-        e = time.time()
-        return 1000 * (e - s) / 100
-
-    simple_time = bench(simple_axpby)
-    custom_time = bench(axpby)
-
-    print(f"Simple axpby: {simple_time:.3f} ms | Custom axpby: {custom_time:.3f} ms")
-
-The results are ``Simple axpby: 1.559 ms | Custom axpby: 0.774 ms``. We see
-modest improvements right away!
-
-This operation is now good to be used to build other operations, in
-:class:`mlx.nn.Module` calls, and also as a part of graph transformations like
-:meth:`grad`.
-
-Scripts
-------
-
-.. admonition:: Download the code
-
-   The full example code is available in `mlx <https://github.com/ml-explore/mlx/tree/main/examples/extensions/>`_.
-
-.. _Accelerate: https://developer.apple.com/documentation/accelerate/blas?language=objc
-.. _Metal: https://developer.apple.com/documentation/metal?language=objc
-.. _Metal-cpp: https://developer.apple.com/metal/cpp/
-.. _`Metal Specification`: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
-.. _`Metal Example`: https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu?language=objc
-.. _nanobind: https://nanobind.readthedocs.io/en/latest/
@@ -1,68 +0,0 @@
-Metal Debugger
-==============
-
-.. currentmodule:: mlx.core
-
-Profiling is a key step for performance optimization. You can build MLX with
-the ``MLX_METAL_DEBUG`` option to improve the Metal debugging and
-optimization workflow. The ``MLX_METAL_DEBUG`` debug option:
-
-* Records source during Metal compilation, for later inspection while
-  debugging.
-* Labels Metal objects such as command queues, improving capture readability.
-
-To build with debugging enabled in Python prepend
-``CMAKE_ARGS="-DMLX_METAL_DEBUG=ON"`` to the build call.
-
-The :func:`metal.start_capture` function initiates a capture of all MLX GPU
-work.
-
-.. note::
-
-   To capture a GPU trace you must run the application with
-   ``MTL_CAPTURE_ENABLED=1``.
-
-.. code-block:: python
-
-    import mlx.core as mx
-
-    a = mx.random.uniform(shape=(512, 512))
-    b = mx.random.uniform(shape=(512, 512))
-    mx.eval(a, b)
-
-    trace_file = "mlx_trace.gputrace"
-
-    # Make sure to run with MTL_CAPTURE_ENABLED=1 and
-    # that the path trace_file does not already exist.
-    mx.metal.start_capture(trace_file)
-
-    for _ in range(10):
-      mx.eval(mx.add(a, b))
-
-    mx.metal.stop_capture()
-
-You can open and replay the GPU trace in Xcode. The ``Dependencies`` view
-has a great overview of all operations. Checkout the `Metal debugger
-documentation`_ for more information.
-
-.. image:: ../_static/metal_debugger/capture.png
-    :class: dark-light
-
-Xcode Workflow
--------------
-
-You can skip saving to a path by running within Xcode. First, generate an
-Xcode project using CMake.
-
-.. code-block::
-
-    mkdir build && cd build
-    cmake .. -DMLX_METAL_DEBUG=ON -G Xcode
-    open mlx.xcodeproj
-
-Select the ``metal_capture`` example schema and run.
-
-.. image:: ../_static/metal_debugger/schema.png
-    :class: dark-light
-
-.. _`Metal debugger documentation`: https://developer.apple.com/documentation/xcode/metal-debugger
@@ -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 17)
-  set(CMAKE_CXX_STANDARD_REQUIRED ON)
-
-
-Depending on how you installed MLX, you may need to tell CMake where to
-find it. 
-
-If you installed MLX with Python, then add the following to the CMake file:
-
-.. code-block:: cmake
-
-  find_package(
-    Python 3.9
-    COMPONENTS Interpreter Development.Module
-    REQUIRED)
-  execute_process(
-    COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE MLX_ROOT)
-
-If you installed the MLX C++ package to a system path, then CMake should be
-able to find it. If you installed it to a non-standard location or CMake can't
-find MLX then set ``MLX_ROOT`` to the location where MLX is installed:
-
-.. code-block:: cmake
-
-  set(MLX_ROOT "/path/to/mlx/")
-
-Next, instruct CMake to find MLX:
-
-.. code-block:: cmake
-
-  find_package(MLX CONFIG REQUIRED)
-
-Finally, add the ``example.cpp`` program as an executable and link MLX.
-
-.. code-block:: cmake
-
-  add_executable(example example.cpp)
-  target_link_libraries(example PRIVATE mlx)
-
-You can build the example with:
-
-.. code-block:: bash
-
-  cmake -B build -DCMAKE_BUILD_TYPE=Release
-  cmake --build build
-
-And run it with:
-
-.. code-block:: bash
-
-  ./build/example
-
-Note ``find_package(MLX CONFIG REQUIRED)`` sets the following variables:
-
-.. list-table:: Package Variables
-   :widths: 20 20 
-   :header-rows: 1
-
-   * - Variable 
-     - Description 
-   * - MLX_FOUND
-     - ``True`` if MLX is found
-   * - MLX_INCLUDE_DIRS
-     - Include directory
-   * - MLX_LIBRARIES
-     - Libraries to link against
-   * - MLX_CXX_FLAGS
-     - Additional compiler flags
-   * - MLX_BUILD_ACCELERATE
-     - ``True`` if MLX was built with Accelerate 
-   * - MLX_BUILD_METAL
-     - ``True`` if MLX was built with Metal
@@ -1,77 +0,0 @@
-.. _linear_regression:
-
-Linear Regression
-----------------
-
-Let's implement a basic linear regression model as a starting point to
-learn MLX. First import the core package and setup some problem metadata:
-
-.. code-block:: python
-
-  import mlx.core as mx
-
-  num_features = 100
-  num_examples = 1_000
-  num_iters = 10_000  # iterations of SGD
-  lr = 0.01  # learning rate for SGD
-
-
-We'll generate a synthetic dataset by:
-
-1. Sampling the design matrix ``X``.
-2. Sampling a ground truth parameter vector ``w_star``.
-3. Compute the dependent values ``y`` by adding Gaussian noise to ``X @ w_star``.
-
-.. code-block:: python
-
-  # True parameters
-  w_star = mx.random.normal((num_features,))
-
-  # Input examples (design matrix)
-  X = mx.random.normal((num_examples, num_features))
-
-  # Noisy labels
-  eps = 1e-2 * mx.random.normal((num_examples,))
-  y = X @ w_star + eps
-
-
-We will use SGD to find the optimal weights. To start, define the squared loss
-and get the gradient function of the loss with respect to the parameters.
-
-.. code-block:: python
-
-  def loss_fn(w):
-      return 0.5 * mx.mean(mx.square(X @ w - y))
-
-  grad_fn = mx.grad(loss_fn)
-
-Start the optimization by initializing the parameters ``w`` randomly. Then
-repeatedly update the parameters for ``num_iters`` iterations. 
-
-.. code-block:: python
-
-  w = 1e-2 * mx.random.normal((num_features,))
-
-  for _ in range(num_iters):
-      grad = grad_fn(w)
-      w = w - lr * grad
-      mx.eval(w)
-
-Finally, compute the loss of the learned parameters and verify that they are
-close to the ground truth parameters.
-
-.. code-block:: python
-
-  loss = loss_fn(w)
-  error_norm = mx.sum(mx.square(w - w_star)).item() ** 0.5
-
-  print(
-      f"Loss {loss.item():.5f}, |w-w*| = {error_norm:.5f}, "
-  )
-  # Should print something close to: Loss 0.00005, |w-w*| = 0.00364
-
-Complete `linear regression
-<https://github.com/ml-explore/mlx/tree/main/examples/python/linear_regression.py>`_
-and `logistic regression
-<https://github.com/ml-explore/mlx/tree/main/examples/python/logistic_regression.py>`_
-examples are available in the MLX GitHub repo.
@@ -1,382 +0,0 @@
-LLM inference
-==============
-
-MLX enables efficient inference of large-ish transformers on Apple silicon
-without compromising on ease of use. In this example we will create an
-inference script for the Llama family of transformer models in which the model
-is defined in less than 200 lines of python.
-
-Implementing the model
----------------------
-
-We will use the neural network building blocks defined in the :mod:`mlx.nn`
-module to concisely define the model architecture. 
-
-Attention layer
-^^^^^^^^^^^^^^^^
-
-We will start with the Llama attention layer which notably uses the RoPE
-positional encoding. [1]_ In addition, our attention layer will optionally use a
-key/value cache that will be concatenated with the provided keys and values to
-support efficient inference.
-
-Our implementation uses :class:`mlx.nn.Linear` for all the projections and
-:class:`mlx.nn.RoPE` for the positional encoding.
-
-.. code-block:: python
-
-    import mlx.core as mx
-    import mlx.nn as nn
-
-    class LlamaAttention(nn.Module):
-        def __init__(self, dims: int, num_heads: int):
-            super().__init__()
-
-            self.num_heads = num_heads
-
-            self.rope = nn.RoPE(dims // num_heads, traditional=True)
-            self.query_proj = nn.Linear(dims, dims, bias=False)
-            self.key_proj = nn.Linear(dims, dims, bias=False)
-            self.value_proj = nn.Linear(dims, dims, bias=False)
-            self.out_proj = nn.Linear(dims, dims, bias=False)
-
-        def __call__(self, queries, keys, values, mask=None, cache=None):
-            queries = self.query_proj(queries)
-            keys = self.key_proj(keys)
-            values = self.value_proj(values)
-
-            # Extract some shapes
-            num_heads = self.num_heads
-            B, L, D = queries.shape
-
-            # Prepare the queries, keys and values for the attention computation
-            queries = queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
-            keys = keys.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
-            values = values.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
-
-            # Add RoPE to the queries and keys and combine them with the cache
-            if cache is not None:
-                key_cache, value_cache = cache
-                queries = self.rope(queries, offset=key_cache.shape[2])
-                keys = self.rope(keys, offset=key_cache.shape[2])
-                keys = mx.concatenate([key_cache, keys], axis=2)
-                values = mx.concatenate([value_cache, values], axis=2)
-            else:
-                queries = self.rope(queries)
-                keys = self.rope(keys)
-
-            # Finally perform the attention computation
-            scale = math.sqrt(1 / queries.shape[-1])
-            scores = (queries * scale) @ keys.transpose(0, 1, 3, 2)
-            if mask is not None:
-                scores = scores + mask
-            scores = mx.softmax(scores, axis=-1)
-            values_hat = (scores @ values).transpose(0, 2, 1, 3).reshape(B, L, -1)
-
-            # Note that we return the keys and values to possibly be used as a cache
-            return self.out_proj(values_hat), (keys, values)
-
-Encoder layer
-^^^^^^^^^^^^^
-
-The other component of the Llama model is the encoder layer which uses RMS
-normalization [2]_ and SwiGLU. [3]_ For RMS normalization we will use
-:class:`mlx.nn.RMSNorm` that is already provided in :mod:`mlx.nn`.
-
-.. code-block:: python
-
-    class LlamaEncoderLayer(nn.Module):
-        def __init__(self, dims: int, mlp_dims: int, num_heads: int):
-            super().__init__()
-
-            self.attention = LlamaAttention(dims, num_heads)
-
-            self.norm1 = nn.RMSNorm(dims)
-            self.norm2 = nn.RMSNorm(dims)
-
-            self.linear1 = nn.Linear(dims, mlp_dims, bias=False)
-            self.linear2 = nn.Linear(dims, mlp_dims, bias=False)
-            self.linear3 = nn.Linear(mlp_dims, dims, bias=False)
-
-        def __call__(self, x, mask=None, cache=None):
-            y = self.norm1(x)
-            y, cache = self.attention(y, y, y, mask, cache)
-            x = x + y
-
-            y = self.norm2(x)
-            a = self.linear1(y)
-            b = self.linear2(y)
-            y = a * mx.sigmoid(a) * b
-            y = self.linear3(y)
-            x = x + y
-
-            return x, cache
-
-Full model
-^^^^^^^^^^
-
-To implement any Llama model we simply have to combine ``LlamaEncoderLayer``
-instances with an :class:`mlx.nn.Embedding` to embed the input tokens.
-
-.. code-block:: python
-
-    class Llama(nn.Module):
-        def __init__(
-            self, num_layers: int, vocab_size: int, dims: int, mlp_dims: int, num_heads: int
-        ):
-            super().__init__()
-
-            self.embedding = nn.Embedding(vocab_size, dims)
-            self.layers = [
-                LlamaEncoderLayer(dims, mlp_dims, num_heads) for _ in range(num_layers)
-            ]
-            self.norm = nn.RMSNorm(dims)
-            self.out_proj = nn.Linear(dims, vocab_size, bias=False)
-
-        def __call__(self, x):
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
-            mask = mask.astype(self.embedding.weight.dtype)
-
-            x = self.embedding(x)
-            for l in self.layers:
-                x, _ = l(x, mask)
-            x = self.norm(x)
-            return self.out_proj(x)
-
-Note that in the implementation above we use a simple list to hold the encoder
-layers but using ``model.parameters()`` will still consider these layers.
-
-Generation
-^^^^^^^^^^^
-
-Our ``Llama`` module can be used for training but not inference as the
-``__call__`` method above processes one input, completely ignores the cache and
-performs no sampling whatsoever. In the rest of this subsection, we will
-implement the inference function as a python generator that processes the
-prompt and then autoregressively yields tokens one at a time.
-
-.. code-block:: python
-
-    class Llama(nn.Module):
-        ...
-
-        def generate(self, x, temp=1.0):
-            cache = []
-
-            # Make an additive causal mask. We will need that to process the prompt.
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
-            mask = mask.astype(self.embedding.weight.dtype)
-
-            # First we process the prompt x the same way as in __call__ but
-            # save the caches in cache
-            x = self.embedding(x)
-            for l in self.layers:
-                x, c = l(x, mask=mask)
-                cache.append(c)  # <--- we store the per layer cache in a
-                                 #      simple python list
-            x = self.norm(x)
-            y = self.out_proj(x[:, -1])  # <--- we only care about the last logits
-                                         #      that generate the next token
-            y = mx.random.categorical(y * (1/temp))
-
-            # y now has size [1]
-            # Since MLX is lazily evaluated nothing is computed yet.
-            # Calling y.item() would force the computation to happen at
-            # this point but we can also choose not to do that and let the
-            # user choose when to start the computation.
-            yield y
-
-            # Now we parsed the prompt and generated the first token we
-            # need to feed it back into the model and loop to generate the
-            # rest.
-            while True:
-                # Unsqueezing the last dimension to add a sequence length
-                # dimension of 1
-                x = y[:, None]
-
-                x = self.embedding(x)
-                for i in range(len(cache)):
-                    # We are overwriting the arrays in the cache list. When
-                    # the computation will happen, MLX will be discarding the
-                    # old cache the moment it is not needed anymore.
-                    x, cache[i] = self.layers[i](x, mask=None, cache=cache[i])
-                x = self.norm(x)
-                y = self.out_proj(x[:, -1])
-                y = mx.random.categorical(y * (1/temp))
-
-                yield y
-
-Putting it all together
-^^^^^^^^^^^^^^^^^^^^^^^
-
-We now have everything we need to create a Llama model and sample tokens from
-it. In the following code, we randomly initialize a small Llama model, process
-6 tokens of prompt and generate 10 tokens.
-
-.. code-block:: python
-
-    model = Llama(num_layers=12, vocab_size=8192, dims=512, mlp_dims=1024, num_heads=8)
-
-    # Since MLX is lazily evaluated nothing has actually been materialized yet.
-    # We could have set the `dims` to 20_000 on a machine with 8GB of RAM and the
-    # code above would still run. Let's actually materialize the model.
-    mx.eval(model.parameters())
-
-    prompt = mx.array([[1, 10, 8, 32, 44, 7]])  # <-- Note the double brackets because we
-                                                #     have a batch dimension even
-                                                #     though it is 1 in this case
-
-    generated = [t for i, t in zip(range(10), model.generate(prompt, 0.8))]
-
-    # Since we haven't evaluated anything, nothing is computed yet. The list
-    # `generated` contains the arrays that hold the computation graph for the
-    # full processing of the prompt and the generation of 10 tokens.
-    #
-    # We can evaluate them one at a time, or all together. Concatenate them or
-    # print them. They would all result in very similar runtimes and give exactly
-    # the same results.
-    mx.eval(generated)
-
-Converting the weights
----------------------
-
-This section assumes that you have access to the original Llama weights and the
-SentencePiece model that comes with them. We will write a small script to
-convert the PyTorch weights to MLX compatible ones and write them in a NPZ file
-that can be loaded directly by MLX.
-
-.. code-block:: python
-
-    import argparse
-    from itertools import starmap
-
-    import numpy as np
-    import torch
-
-    def map_torch_to_mlx(key, value):
-        if "tok_embedding" in key:
-            key = "embedding.weight"
-
-        elif "norm" in key:
-            key = key.replace("attention_norm", "norm1").replace("ffn_norm", "norm2")
-
-        elif "wq" in key or "wk" in key or "wv" in key or "wo" in key:
-            key = key.replace("wq", "query_proj")
-            key = key.replace("wk", "key_proj")
-            key = key.replace("wv", "value_proj")
-            key = key.replace("wo", "out_proj")
-
-        elif "w1" in key or "w2" in key or "w3" in key:
-            # The FFN is a separate submodule in PyTorch
-            key = key.replace("feed_forward.w1", "linear1")
-            key = key.replace("feed_forward.w3", "linear2")
-            key = key.replace("feed_forward.w2", "linear3")
-
-        elif "output" in key:
-            key = key.replace("output", "out_proj")
-
-        elif "rope" in key:
-            return None, None
-
-        return key, value.numpy()
-
-
-    if __name__ == "__main__":
-        parser = argparse.ArgumentParser(description="Convert Llama weights to MLX")
-        parser.add_argument("torch_weights")
-        parser.add_argument("output_file")
-        args = parser.parse_args()
-
-        state = torch.load(args.torch_weights)
-        np.savez(
-            args.output_file,
-            **{k: v for k, v in starmap(map_torch_to_mlx, state.items()) if k is not None}
-        )
-
-
-Weight loading and benchmarking
-------------------------------
-
-After converting the weights to be compatible to our implementation, all that is
-left is to load them from disk and we can finally use the LLM to generate text.
-We can load numpy format files using the :func:`mlx.core.load` operation.
-
-To create a parameter dictionary from the key/value representation of NPZ files
-we will use the :func:`mlx.utils.tree_unflatten` helper method as follows:
-
-.. code-block:: python
-
-    from mlx.utils import tree_unflatten
-
-    model.update(tree_unflatten(list(mx.load(weight_file).items())))
-
-:meth:`mlx.utils.tree_unflatten` will take keys from the NPZ file that look
-like ``layers.2.attention.query_proj.weight`` and will transform them to
-
-.. code-block:: python
-
-   {"layers": [..., ..., {"attention": {"query_proj": {"weight": ...}}}]}
-
-which can then be used to update the model. Note that the method above incurs
-several unnecessary copies from disk to numpy and then from numpy to MLX. It
-will be replaced in the future with direct loading to MLX.
-
-You can download the full example code in `mlx-examples`_. Assuming, the
-existence of ``weights.pth`` and ``tokenizer.model`` in the current working
-directory we can play around with our inference script as follows (the timings
-are representative of an M1 Ultra and the 7B parameter Llama model):
-
-.. code-block:: bash
-
-    $ python convert.py weights.pth llama-7B.mlx.npz
-    $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely'
-    [INFO] Loading model from disk: 5.247 s
-    Press enter to start generation
-    ------
-    , having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down,
-    ------
-    [INFO] Prompt processing: 0.437 s
-    [INFO] Full generation: 4.330 s
-
-We observe that 4.3 seconds are required to generate 100 tokens and 0.4 seconds
-of those are spent processing the prompt. This amounts to a little over **39 ms
-per token**.
-
-By running with a much bigger prompt we can see that the per token generation
-time as well as the prompt processing time remains almost constant.
-
-.. code-block:: bash
-
-    $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not'
-    [INFO] Loading model from disk: 5.247 s
-    Press enter to start generation
-    ------
-    take his eyes from the ground. “What is it you are waiting for?” said I. “I am not accustomed to be thus questioned,” said he. “You look like a reasonable man—tell me, then, what are you waiting for?” “You would not understand,” he replied; “and how could you help me, if I were to tell you?” “I should not only understand, but would do all that I could,” said I. He did not
-    ------
-    [INFO] Prompt processing: 0.579 s
-    [INFO] Full generation: 4.690 s
-    $ python llama.py --num-tokens 500 llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not'
-    [INFO] Loading model from disk: 5.628 s
-    Press enter to start generation
-    ------
-    take his eyes from the ground. “What is it you are waiting for?” said I. “I am not accustomed to be thus questioned,” said he. “You look like a reasonable man—tell me, then, what are you waiting for?” “You would not understand,” he replied; “and how could you help me, if I were to tell you?” “I should not only understand, but would do all that I could,” said I. He did not reply, but still went on looking at the ground, and took hold of his bundle with a nervous trembling. I waited some time, and then resumed. “It is of no use to say you would not understand, if I were to tell you,” said he. “I have not told you why I am waiting for him,” said I. “And I am sure I should not understand,” replied he. “I will tell you then,” said I, “and, perhaps, you would not be surprised.” “No matter,” said he, “I shall be surprised anyhow; so tell me why you are waiting for him.” “He is my friend,” said I. “Yes,” said he, with a slight smile, “I know.” “He has been kind to me,” said I, “and I am waiting for him. I want to see him, and could have waited as I am now, for a much longer time.” “He will not soon come,” said he. “Unless he sees you here, he will not know of your having waited, and he will be very unlikely to come.” “No matter,” said I, “I shall wait for him.” “This is a strange thing,” said he, still with the same amused smile. “How did you know,” said I, “that he was coming? How should you be waiting?” “That is my secret,” said he. “And you expect him?” “Yes,” said I. “Are you disappointed then, if he does not come?” “No,” said I, “it is his secret, not mine.” “If he comes,” said he, “do you mean to go straight away?” “Yes,” said I, “I cannot be happy if I do not go straight away after him.” “Did you know this place before?” asked he. “Yes,” said I. “Is there any shop to buy food here?” “
-    ------
-    [INFO] Prompt processing: 0.633 s
-    [INFO] Full generation: 21.475 s
-
-Scripts
-------
-
-.. admonition:: Download the code
-
-   The full example code is available in `mlx-examples`_.
-
-.. _mlx-examples: https://github.com/ml-explore/mlx-examples/tree/main/llms/llama
-
-.. [1] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B. and Liu, Y., 2021.
-   Roformer: Enhanced transformer with rotary position embedding. arXiv
-   preprint arXiv:2104.09864.
-.. [2] Zhang, B. and Sennrich, R., 2019. Root mean square layer normalization.
-   Advances in Neural Information Processing Systems, 32.
-.. [3] Shazeer, N., 2020. Glu variants improve transformer. arXiv preprint
-   arXiv:2002.05202.
@@ -1,134 +0,0 @@
-.. _mlp:
-
-Multi-Layer Perceptron
----------------------
-
-In this example we'll learn to use ``mlx.nn`` by implementing a simple
-multi-layer perceptron to classify MNIST.
-
-As a first step import the MLX packages we need:
-
-.. code-block:: python
-
-  import mlx.core as mx
-  import mlx.nn as nn
-  import mlx.optimizers as optim
-
-  import numpy as np
-
-
-The model is defined as the ``MLP`` class which inherits from
-:class:`mlx.nn.Module`. We follow the standard idiom to make a new module:
-
-1. Define an ``__init__`` where the parameters and/or submodules are setup. See
-   the :ref:`Module class docs<module_class>` for more information on how
-   :class:`mlx.nn.Module` registers parameters.
-2. Define a ``__call__`` where the computation is implemented.
-
-.. code-block:: python
-
-  class MLP(nn.Module):
-      def __init__(
-          self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
-      ):
-          super().__init__()
-          layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
-          self.layers = [
-              nn.Linear(idim, odim)
-              for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
-          ]
-
-      def __call__(self, x):
-          for l in self.layers[:-1]:
-              x = mx.maximum(l(x), 0.0)
-          return self.layers[-1](x)
-
-
-We define the loss function which takes the mean of the per-example cross
-entropy loss.  The ``mlx.nn.losses`` sub-package has implementations of some
-commonly used loss functions.
-
-.. code-block:: python
-
-  def loss_fn(model, X, y):
-      return mx.mean(nn.losses.cross_entropy(model(X), y))
-
-We also need a function to compute the accuracy of the model on the validation
-set:
-
-.. code-block:: python
-
-  def eval_fn(model, X, y):
-      return mx.mean(mx.argmax(model(X), axis=1) == y)
-
-Next, setup the problem parameters and load the data. To load the data, you need our
-`mnist data loader
-<https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py>`_, which
-we will import as ``mnist``.
-
-.. code-block:: python
-
-  num_layers = 2
-  hidden_dim = 32
-  num_classes = 10
-  batch_size = 256
-  num_epochs = 10
-  learning_rate = 1e-1
-
-  # Load the data
-  import mnist 
-  train_images, train_labels, test_images, test_labels = map(
-      mx.array, mnist.mnist()
-  )
-
-Since we're using SGD, we need an iterator which shuffles and constructs
-minibatches of examples in the training set:
-
-.. code-block:: python
-
-  def batch_iterate(batch_size, X, y):
-      perm = mx.array(np.random.permutation(y.size))
-      for s in range(0, y.size, batch_size):
-          ids = perm[s : s + batch_size]
-          yield X[ids], y[ids]
-
-
-Finally, we put it all together by instantiating the model, the
-:class:`mlx.optimizers.SGD` optimizer, and running the training loop:
-
-.. code-block:: python
-
-  # Load the model
-  model = MLP(num_layers, train_images.shape[-1], hidden_dim, num_classes)
-  mx.eval(model.parameters())
-
-  # Get a function which gives the loss and gradient of the
-  # loss with respect to the model's trainable parameters
-  loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
-
-  # Instantiate the optimizer
-  optimizer = optim.SGD(learning_rate=learning_rate)
-
-  for e in range(num_epochs):
-      for X, y in batch_iterate(batch_size, train_images, train_labels):
-          loss, grads = loss_and_grad_fn(model, X, y)
-
-          # Update the optimizer state and model parameters
-          # in a single call
-          optimizer.update(model, grads)
-
-          # Force a graph evaluation
-          mx.eval(model.parameters(), optimizer.state)
-
-      accuracy = eval_fn(model, test_images, test_labels)
-      print(f"Epoch {e}: Test accuracy {accuracy.item():.3f}")
-
-
-.. note::
-  The :func:`mlx.nn.value_and_grad` function is a convenience function to get
-  the gradient of a loss with respect to the trainable parameters of a model.
-  This should not be confused with :func:`mlx.core.value_and_grad`.
-
-The model should train to a decent accuracy (about 95%) after just a few passes
-over the training set. The `full example <https://github.com/ml-explore/mlx-examples/tree/main/mnist>`_
-is available in the MLX GitHub repo.
@@ -1,93 +0,0 @@
-MLX
-===
-
-MLX is a NumPy-like array framework designed for efficient and flexible machine
-learning on Apple silicon, brought to you by Apple machine learning research.
-
-The Python API closely follows NumPy with a few exceptions. MLX also has a
-fully featured C++ API which closely follows the Python API.
-
-The main differences between MLX and NumPy are:
-
- - **Composable function transformations**: MLX has 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.
- - **Multi-device**: Operations can run on any of the supported devices (CPU,
-   GPU, ...)
-
-The design of MLX is inspired by frameworks like `PyTorch
-<https://pytorch.org/>`_, `Jax <https://github.com/google/jax>`_, and
-`ArrayFire <https://arrayfire.org/>`_. A notable difference from these
-frameworks and MLX 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 performing data copies. Currently supported device types
-are the CPU and GPU.
-
-.. toctree::
-   :caption: Install
-   :maxdepth: 1
-
-   install
-
-.. toctree::
-   :caption: Usage 
-   :maxdepth: 1
-
-   usage/quick_start
-   usage/lazy_evaluation
-   usage/unified_memory
-   usage/indexing
-   usage/saving_and_loading
-   usage/function_transforms
-   usage/compile
-   usage/numpy
-   usage/distributed
-   usage/using_streams
-   usage/export
-
-.. toctree::
-   :caption: Examples
-   :maxdepth: 1
-
-   examples/linear_regression
-   examples/mlp
-   examples/llama-inference
-
-.. toctree::
-   :caption: Python API Reference
-   :maxdepth: 1
-
-   python/array
-   python/data_types
-   python/devices_and_streams
-   python/export
-   python/ops
-   python/random
-   python/transforms
-   python/fast
-   python/fft
-   python/linalg
-   python/metal
-   python/cuda
-   python/memory_management
-   python/nn
-   python/optimizers
-   python/distributed
-   python/tree_utils
-
-.. toctree::
-   :caption: C++ API Reference
-   :maxdepth: 1
-
-   cpp/ops
-
-.. toctree::
-   :caption: Further Reading
-   :maxdepth: 1
-
-   dev/extensions
-   dev/metal_debugger
-   dev/custom_metal_kernels
-   dev/mlx_in_cpp
@@ -1,345 +0,0 @@
-.. _build_and_install:
-
-Build and Install
-=================
-
-Python Installation
-------------------
-
-MLX is available on PyPI. All you have to do to use MLX with your own Apple
-silicon computer is
-
-.. code-block:: shell
-
-    pip install mlx
-
-To install from PyPI your system must meet the following requirements:
-
- 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 >= 13.5
-    It is highly recommended to use macOS 14 (Sonoma)
-
-CUDA
-^^^^
-
-MLX has a CUDA backend which you can install with:
-
-.. code-block:: shell
-
-    pip install mlx[cuda]
-
-To install the CUDA package from PyPi your system must meet the following
-requirements:
-
- Nvidia architecture >= SM 7.0 (Volta)
- Nvidia driver >= 550.54.14
- CUDA toolkit >= 12.0
- Linux distribution with glibc >= 2.35
- Python >= 3.9
-
-
-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.9
-
-
-Troubleshooting
-^^^^^^^^^^^^^^^
-
-*My OS and Python versions are in the required range but pip still does not find
-a matching distribution.*
-
-Probably you are using a non-native Python. The output of
-
-.. code-block:: shell
-
-  python -c "import platform; print(platform.processor())"
-
-should be ``arm``. If it is ``i386`` (and you have M series machine) then you
-are using a non-native Python. Switch your Python to a native Python. A good
-way to do this is with `Conda <https://stackoverflow.com/q/65415996>`_.
-
-
-Build from source
-----------------
-
-Build Requirements
-^^^^^^^^^^^^^^^^^^
-
- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
- `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
- Xcode >= 15.0 and macOS SDK >= 14.0
-
-.. note::
-   Ensure your shell environment is native ``arm``, not ``x86`` via Rosetta. If
-   the output of ``uname -p`` is ``x86``, see the :ref:`troubleshooting section <build shell>` below.
-
-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>`_:
-
-.. code-block:: shell
-
-   git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
-
-Then simply build and install MLX using pip:
-
-.. code-block:: shell
-
-  pip install .
-
-For developing, install the package with development dependencies, and use an
-editable install:
-
-.. code-block:: shell
-
-  pip install -e ".[dev]"
-
-Once the development dependencies are installed, you can build faster with:
-
-.. code-block:: shell
-
- python setup.py build_ext --inplace
-
-Run the tests with:
-
-.. code-block:: shell
-
-  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
-by cloning MLX from `its GitHub repo
-<https://github.com/ml-explore/mlx>`_:
-
-.. code-block:: shell
-
-   git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
-
-Create a build directory and run CMake and make:
-
-.. code-block:: shell
-
-   mkdir -p build && cd build
-   cmake .. && make -j
-
-Run tests with:
-
-.. code-block:: shell
-
-   make test
-
-Install with:
-
-.. code-block:: shell
-
-   make install
-
-Note that the built ``mlx.metallib`` file should be either at the same
-directory as the executable statically linked to ``libmlx.a`` or the
-preprocessor constant ``METAL_PATH`` should be defined at build time and it
-should point to the path to the built metal library.
-
-.. list-table:: Build Options
-   :widths: 25 8
-   :header-rows: 1
-
-   * - Option
-     - Default
-   * - MLX_BUILD_TESTS
-     - ON
-   * - MLX_BUILD_EXAMPLES
-     - OFF
-   * - MLX_BUILD_BENCHMARKS
-     - OFF
-   * - MLX_BUILD_METAL
-     - ON
-   * - MLX_BUILD_CPU
-     - ON
-   * - MLX_BUILD_PYTHON_BINDINGS
-     - OFF
-   * - MLX_METAL_DEBUG
-     - OFF
-   * - MLX_BUILD_SAFETENSORS
-     - ON
-   * - MLX_BUILD_GGUF
-     - ON
-   * - MLX_METAL_JIT
-     - OFF
-
-.. note::
-
-    If you have multiple Xcode installations and wish to use
-    a specific one while building, you can do so by adding the
-    following environment variable before building
-
-    .. code-block:: shell
-
-      export DEVELOPER_DIR="/path/to/Xcode.app/Contents/Developer/"
-
-    Further, you can use the following command to find out which
-    macOS SDK will be used
-
-    .. code-block:: shell
-
-      xcrun -sdk macosx --show-sdk-version
-
-
-Binary Size Minimization
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-To produce a smaller binary use the CMake flags ``CMAKE_BUILD_TYPE=MinSizeRel``
-and ``BUILD_SHARED_LIBS=ON``.
-
-The MLX CMake build has several additional options to make smaller binaries.
-For example, if you don't need the CPU backend or support for safetensors and
-GGUF, you can do:
-
-.. code-block:: shell
-
-  cmake .. \
-    -DCMAKE_BUILD_TYPE=MinSizeRel \
-    -DBUILD_SHARED_LIBS=ON \
-    -DMLX_BUILD_CPU=OFF \
-    -DMLX_BUILD_SAFETENSORS=OFF \
-    -DMLX_BUILD_GGUF=OFF \
-    -DMLX_METAL_JIT=ON
-
-THE ``MLX_METAL_JIT`` flag minimizes the size of the MLX Metal library which
-contains pre-built GPU kernels. This substantially reduces the size of the
-Metal library by run-time compiling kernels the first time they are used in MLX
-on a given machine. Note run-time compilation incurs a cold-start cost which can
-be anwywhere from a few hundred millisecond to a few seconds depending on the
-application. Once a kernel is compiled, it will be cached by the system. The
-Metal kernel cache persists 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
-^^^^^^^^^^^^^^^
-
-Metal not found
-~~~~~~~~~~~~~~~
-
-You see the following error when you try to build:
-
-.. code-block:: shell
-
-  error: unable to find utility "metal", not a developer tool or in PATH
-
-To fix this, first make sure you have Xcode installed:
-
-.. code-block:: shell
-
-  xcode-select --install
-
-Then set the active developer directory:
-
-.. code-block:: shell
-
-  sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
-
-x86 Shell
-~~~~~~~~~
-
-.. _build shell:
-
-If the output of ``uname -p``  is ``x86`` then your shell is running as x86 via
-Rosetta instead of natively.
-
-To fix this, find the application in Finder (``/Applications`` for iTerm,
-``/Applications/Utilities`` for Terminal), right-click, and click “Get Info”.
-Uncheck “Open using Rosetta”, close the “Get Info” window, and restart your
-terminal.
-
-Verify the terminal is now running natively the following command:
-
-.. code-block:: shell
-
-  $ uname -p
-  arm
-
-Also check that cmake is using the correct architecture:
-
-.. code-block:: shell
-
-  $ cmake --system-information | grep CMAKE_HOST_SYSTEM_PROCESSOR
-  CMAKE_HOST_SYSTEM_PROCESSOR "arm64"
-
-If you see ``"x86_64"``, try re-installing ``cmake``. If you see ``"arm64"``
-but the build errors out with "Building for x86_64 on macOS is not supported."
-wipe your build cache with ``rm -rf build/`` and try again.
@@ -1,28 +0,0 @@
-mlx.core.Device
-===============
-
-.. currentmodule:: mlx.core
-
-.. autoclass:: Device
-
-   
-   .. automethod:: __init__
-
-   
-   .. rubric:: Methods
-
-   .. autosummary::
-   
-      ~Device.__init__
-   
-   
-
-   
-   
-   .. rubric:: Attributes
-
-   .. autosummary::
-   
-      ~Device.type
-   
-   
@@ -1,28 +0,0 @@
-mlx.core.Dtype
-==============
-
-.. currentmodule:: mlx.core
-
-.. autoclass:: Dtype
-
-   
-   .. automethod:: __init__
-
-   
-   .. rubric:: Methods
-
-   .. autosummary::
-   
-      ~Dtype.__init__
-   
-   
-
-   
-   
-   .. rubric:: Attributes
-
-   .. autosummary::
-   
-      ~Dtype.size
-   
-   
@@ -1,29 +0,0 @@
-mlx.core.DtypeCategory
-======================
-
-.. currentmodule:: mlx.core
-
-.. autoclass:: DtypeCategory
-
-   
-   .. automethod:: __init__
-
-   
-   
-
-   
-   
-   .. rubric:: Attributes
-
-   .. autosummary::
-   
-      ~DtypeCategory.complexfloating
-      ~DtypeCategory.floating
-      ~DtypeCategory.inexact
-      ~DtypeCategory.signedinteger
-      ~DtypeCategory.unsignedinteger
-      ~DtypeCategory.integer
-      ~DtypeCategory.number
-      ~DtypeCategory.generic
-   
-   
@@ -1,6 +0,0 @@
-mlx.core.abs
-============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: abs
@@ -1,6 +0,0 @@
-mlx.core.add
-============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: add
@@ -1,6 +0,0 @@
-mlx.core.addmm
-==============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: addmm
@@ -1,6 +0,0 @@
-mlx.core.all
-============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: all
@@ -1,6 +0,0 @@
-mlx.core.allclose
-=================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: allclose
@@ -1,6 +0,0 @@
-mlx.core.any
-============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: any
@@ -1,6 +0,0 @@
-mlx.core.arange
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arange
@@ -1,6 +0,0 @@
-mlx.core.arccos
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arccos
@@ -1,6 +0,0 @@
-mlx.core.arccosh
-================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arccosh
@@ -1,6 +0,0 @@
-mlx.core.arcsin
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arcsin
@@ -1,6 +0,0 @@
-mlx.core.arcsinh
-================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arcsinh
@@ -1,6 +0,0 @@
-mlx.core.arctan
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arctan
@@ -1,6 +0,0 @@
-mlx.core.arctan2
-================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arctan2
@@ -1,6 +0,0 @@
-mlx.core.arctanh
-================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: arctanh
@@ -1,6 +0,0 @@
-mlx.core.argmax
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: argmax
@@ -1,6 +0,0 @@
-mlx.core.argmin
-===============
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: argmin
@@ -1,6 +0,0 @@
-mlx.core.argpartition
-=====================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: argpartition
@@ -1,6 +0,0 @@
-mlx.core.argsort
-================
-
-.. currentmodule:: mlx.core
-
-.. autofunction:: argsort
@@ -1,6 +0,0 @@
-mlx.core.array.T
-================
-
-.. currentmodule:: mlx.core
-
-.. autoproperty:: array.T
@@ -1,6 +0,0 @@
-mlx.core.array.abs
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.abs
@@ -1,6 +0,0 @@
-mlx.core.array.all
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.all
@@ -1,6 +0,0 @@
-mlx.core.array.any
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.any
@@ -1,6 +0,0 @@
-mlx.core.array.argmax
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.argmax
@@ -1,6 +0,0 @@
-mlx.core.array.argmin
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.argmin
@@ -1,6 +0,0 @@
-mlx.core.array.astype
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.astype
@@ -1,6 +0,0 @@
-mlx.core.array.at
-=================
-
-.. currentmodule:: mlx.core
-
-.. autoproperty:: array.at
@@ -1,6 +0,0 @@
-mlx.core.array.conj
-===================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.conj
@@ -1,6 +0,0 @@
-mlx.core.array.cos
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.cos
@@ -1,6 +0,0 @@
-mlx.core.array.cummax
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.cummax
@@ -1,6 +0,0 @@
-mlx.core.array.cummin
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.cummin
@@ -1,6 +0,0 @@
-mlx.core.array.cumprod
-======================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.cumprod
@@ -1,6 +0,0 @@
-mlx.core.array.cumsum
-=====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.cumsum
@@ -1,6 +0,0 @@
-mlx.core.array.diag
-===================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.diag
@@ -1,6 +0,0 @@
-mlx.core.array.diagonal
-=======================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.diagonal
@@ -1,6 +0,0 @@
-mlx.core.array.dtype
-====================
-
-.. currentmodule:: mlx.core
-
-.. autoproperty:: array.dtype
@@ -1,6 +0,0 @@
-mlx.core.array.exp
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.exp
@@ -1,6 +0,0 @@
-mlx.core.array.flatten
-======================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.flatten
@@ -1,6 +0,0 @@
-mlx.core.array.imag
-===================
-
-.. currentmodule:: mlx.core
-
-.. autoproperty:: array.imag
@@ -1,6 +0,0 @@
-mlx.core.array.item
-===================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.item
@@ -1,6 +0,0 @@
-mlx.core.array.itemsize
-=======================
-
-.. currentmodule:: mlx.core
-
-.. autoproperty:: array.itemsize
@@ -1,6 +0,0 @@
-mlx.core.array.log
-==================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.log
@@ -1,6 +0,0 @@
-mlx.core.array.log10
-====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.log10
@@ -1,6 +0,0 @@
-mlx.core.array.log1p
-====================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.log1p
@@ -1,6 +0,0 @@
-mlx.core.array.log2
-===================
-
-.. currentmodule:: mlx.core
-
-.. automethod:: array.log2
--- a/Show More
+++ b/Show More