fix test

.circleci/config.yml

@@ -0,0 +1,413 @@
+version: 2.1
+
+orbs:
+  apple: ml-explore/pr-approval@0.1.0
+
+parameters:
+  nightly_build:
+    type: boolean
+    default: false
+  weekly_build:
+    type: boolean
+    default: false
+  test_release:
+    type: boolean
+    default: false
+  linux_release:
+    type: boolean
+    default: false
+
+jobs:
+  build_documentation:
+    parameters:
+      upload-docs:
+        type: boolean
+        default: false
+    macos:
+      xcode: "15.2.0"
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install
+          command: |
+            brew install python@3.9
+            brew install doxygen
+            python3.9 -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install -r docs/requirements.txt
+            CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` pip install . -v
+      - when:
+          condition:
+            not: << parameters.upload-docs >>
+          steps:
+            - run:
+               name: Build documentation
+               command: |
+                 source env/bin/activate
+                 cd docs && doxygen && make html O=-W
+      - when:
+          condition: << parameters.upload-docs >>
+          steps:
+            - add_ssh_keys:
+                fingerprints:
+                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
+            - run:
+               name: Upload documentation
+               command: |
+                 source env/bin/activate
+                 git config user.email "mlx@group.apple.com"
+                 git config user.name "CircleCI Docs"
+                 git checkout gh-pages
+                 git rebase main
+                 cd docs
+                 git rm -rf build/html
+                 doxygen && make html O=-W
+                 git add -f build/html
+                 git commit -m "rebase"
+                 git push -f origin gh-pages
+
+  linux_build_and_test:
+    docker:
+      - image: cimg/python:3.9
+
+    steps:
+      - checkout
+      - run:
+          name: Run style checks
+          command: |
+            pip install pre-commit
+            pre-commit run --all
+            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
+      - run:
+          name: Install dependencies
+          command: |
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install numpy
+            sudo apt-get update
+            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
+      - run:
+          name: Install Python package
+          command: |
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              python3 setup.py build_ext --inplace
+            CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              python3 setup.py develop
+      - run:
+          name: Generate package stubs
+          command: |
+            echo "stubs"
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Run Python tests
+          command: |
+            python3 -m unittest discover python/tests -v
+      - run:
+          name: Build CPP only
+          command: |
+            mkdir -p build && cd build 
+            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
+            make -j `nproc`
+      - run:
+          name: Run CPP tests
+          command: ./build/tests/tests
+
+  mac_build_and_test:
+    parameters:
+      xcode_version:
+        type: string
+        default: "15.2.0"
+    macos:
+      xcode: << parameters.xcode_version >>
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            brew install python@3.9
+            brew install openmpi
+            python3.9 -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install numpy
+            pip install torch
+            pip install tensorflow
+            pip install unittest-xml-reporting
+      - run:
+          name: Install Python package
+          command: |
+            source env/bin/activate
+            DEBUG=1 CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` pip install -e . -v
+      - run:
+          name: Generate package stubs
+          command: |
+            source env/bin/activate
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Run Python tests
+          command: |
+            source env/bin/activate
+            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
+            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
+            mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
+      - run:
+          name: Build example extension
+          command: |
+            source env/bin/activate
+            cd examples/extensions
+            pip install -r requirements.txt
+            python setup.py build_ext -j8
+      - store_test_results:
+          path: test-results
+      - run:
+          name: Build CPP only
+          command: |
+            source env/bin/activate
+            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
+      - run:
+          name: Run CPP tests
+          command: |
+            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
+      - run:
+          name: Build small binary
+          command: |
+            source env/bin/activate
+            cd build/
+            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
+              -DBUILD_SHARED_LIBS=ON \
+              -DMLX_BUILD_CPU=OFF \
+              -DMLX_BUILD_SAFETENSORS=OFF \
+              -DMLX_BUILD_GGUF=OFF \
+              -DMLX_METAL_JIT=ON
+            make -j `sysctl -n hw.ncpu`
+      - run:
+          name: Run Python tests with JIT
+          command: |
+            source env/bin/activate
+            CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
+              pip install -e . -v
+            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
+              METAL_DEBUG_ERROR_MODE=0 \
+              python -m xmlrunner discover -v python/tests -o test-results/gpu_jit
+
+  build_release:
+    parameters:
+      python_version:
+        type: string
+        default: "3.9"
+      xcode_version:
+        type: string
+        default: "15.2.0"
+      build_env:
+        type: string
+        default: ""
+    macos:
+      xcode: << parameters.xcode_version >>
+    resource_class: macos.m1.medium.gen1
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            brew install python@<< parameters.python_version >>
+            brew install openmpi
+            python<< parameters.python_version >> -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install --upgrade setuptools
+            pip install numpy
+            pip install twine
+            pip install build
+      - run:
+          name: Install Python package
+          command: |
+            source env/bin/activate
+            DEV_RELEASE=1 \
+              CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              pip install . -v
+      - run:
+          name: Generate package stubs
+          command: |
+            source env/bin/activate
+            pip install typing_extensions
+            python setup.py generate_stubs 
+      - run:
+          name: Build Python package
+          command: |
+            source env/bin/activate
+            << parameters.build_env >> \
+              CMAKE_BUILD_PARALLEL_LEVEL=`sysctl -n hw.ncpu` \
+              python -m build -w
+      - when:
+          condition: << parameters.build_env >>
+          steps:
+            - run:
+                name: Upload package
+                command: |
+                  source env/bin/activate
+                  twine upload dist/*
+      - store_artifacts:
+          path: dist/
+
+  build_linux_release:
+    parameters:
+      python_version:
+        type: string
+        default: "3.9"
+      extra_env:
+        type: string
+        default: "DEV_RELEASE=1"
+    docker:
+      - image: ubuntu:20.04
+    steps:
+      - checkout
+      - run:
+          name: Build wheel
+          command: |
+            PYTHON=python<< parameters.python_version >>
+            apt-get update
+            apt-get upgrade -y
+            DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
+            apt-get install -y apt-utils
+            apt-get install -y software-properties-common
+            add-apt-repository -y ppa:deadsnakes/ppa
+            apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
+            apt-get install -y libblas-dev liblapack-dev liblapacke-dev
+            apt-get install -y build-essential git
+            $PYTHON -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install --upgrade cmake
+            pip install nanobind==2.2.0
+            pip install --upgrade setuptools
+            pip install numpy
+            pip install auditwheel
+            pip install patchelf
+            pip install build
+            pip install twine
+            << parameters.extra_env >> \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              pip install . -v
+            pip install typing_extensions
+            python setup.py generate_stubs 
+            << parameters.extra_env >> \
+              CMAKE_BUILD_PARALLEL_LEVEL=`nproc` \
+              python -m build --wheel
+            auditwheel show dist/*
+            auditwheel repair dist/* --plat manylinux_2_31_x86_64
+      - run:
+          name: Upload package
+          command: |
+            source env/bin/activate
+            twine upload wheelhouse/*
+      - store_artifacts:
+          path: wheelhouse/
+
+workflows:
+  build_and_test:
+    when:
+      and:
+        - matches:
+            pattern: "^(?!pull/)[-\\w]+$"
+            value: << pipeline.git.branch >>
+        - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.weekly_build >>
+        - not: << pipeline.parameters.test_release >>
+    jobs:
+      - mac_build_and_test:
+          matrix:
+            parameters:
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
+      - linux_build_and_test
+      - build_documentation 
+
+  build_pypi_release:
+    when:
+      and:
+        - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.weekly_build >>
+        - not: << pipeline.parameters.test_release >>
+    jobs:
+      - build_release:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              xcode_version: ["15.0.0", "15.2.0"]
+              build_env: ["PYPI_RELEASE=1"]
+      - build_documentation:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+          upload-docs: true
+
+  prb:
+    when:
+      matches:
+        pattern: "^pull/\\d+(/head)?$"
+        value: << pipeline.git.branch >>
+    jobs:
+      - hold:
+          type: approval
+      - apple/authenticate:
+          context: pr-approval
+      - mac_build_and_test:
+          requires: [ hold ]
+          matrix:
+            parameters:
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
+      - linux_build_and_test:
+          requires: [ hold ]
+  nightly_build:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.nightly_build >>
+    jobs:
+      - build_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              xcode_version: ["15.0.0", "15.2.0"]
+  weekly_build:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.weekly_build >>
+    jobs:
+      - build_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              xcode_version: ["15.0.0", "15.2.0", "16.0.0"]
+              build_env: ["DEV_RELEASE=1"]
+  linux_test_release:
+    when:
+      and:
+        - equal: [ main, << pipeline.git.branch >> ]
+        - << pipeline.parameters.linux_release >>
+    jobs:
+      - build_linux_release:
+          matrix:
+            parameters:
+              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+              extra_env: ["PYPI_RELEASE=1"]

.github/actions/build-cuda-release/action.yml

@@ -1,15 +0,0 @@
-name: 'Build CUDA wheel'
-description: 'Build CUDA wheel'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build package
-      shell: bash
-      env:
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
-      run: |
-        pip install auditwheel build patchelf setuptools
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        bash python/scripts/repair_cuda.sh

.github/actions/build-docs/action.yml

@@ -1,38 +0,0 @@
-name: 'Build Documentation'
-description: 'Build documentation'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Setup machine
-      uses: ./.github/actions/setup-linux
-
-    - name: Install dependencies
-      shell: bash
-      run: |
-        sudo apt-get install -y doxygen
-        source .venv/bin/activate
-        pip install -r docs/requirements.txt
-        pip install . -v
-  
-    - name: Build documentation
-      shell: bash
-      run: |
-        source .venv/bin/activate
-        cd docs
-        doxygen
-        make html O=-W
-    
-    - name: Create artifact tar
-      shell: bash
-      run: tar -cf artifact.tar -C docs --dereference build/html index.html
-
-    # Do it manually because upload-pages-artifact requires gtar
-    - name: Upload artifact
-      id: upload-artifact
-      uses: actions/upload-artifact@v5
-      with:
-        name: github-pages
-        path: artifact.tar
-        retention-days: 1
-        if-no-files-found: error

.github/actions/build-linux-release/action.yml

@@ -1,40 +0,0 @@
-name: 'Build Linux wheel'
-description: 'Build Linux wheel'
-
-inputs:
-  build-backend:
-    description: 'Build the backend mlx-cpu package'
-    type: boolean
-    required: false
-    default: false
-  arch:
-    description: 'Platform architecture tag'
-    required: true
-    type: choice
-    options:
-      - x86_64
-      - aarch64
-
-runs:
-  using: "composite"
-  steps:
-    - name: Generate package stubs
-      shell: bash
-      run: |
-        pip install -e ".[dev]" -v
-        pip install typing_extensions
-        python setup.py generate_stubs
-    - name: Build Python package
-      shell: bash
-      run: |
-        pip install auditwheel patchelf build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        bash python/scripts/repair_linux.sh ${{ inputs.arch }}
-    - 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 }}

.github/actions/build-linux/action.yml

@@ -1,41 +0,0 @@
-name: 'Build and Test on Linux'
-
-inputs:
-  toolkit:
-    description: 'The toolkit to build with'
-    required: false
-    default: 'cpu'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      id: python_build
-      shell: sh
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: >-
-          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
-          -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }}
-      run: |
-        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
-          # There is no GPU in arm64 runner, use a common arch.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
-          # Can not build tests when the built executables can not run.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
-        fi
-        pip install --no-build-isolation -e ".[dev]" -v
-        # Pass the CMAKE_ARGS to following steps.
-        echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
-
-    - name: Generate package stubs
-      shell: sh
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
-    - 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)

.github/actions/build-macos-release/action.yml

@@ -1,34 +0,0 @@
-name: 'Build macOS release'
-description: 'Build MLX releases macOS'
-
-inputs:
-  macos-target:
-    description: 'macOS build target'
-    required: false
-    default: '15.0'
-  build-backend:
-    description: 'Build the backend mlx-metal package'
-    type: boolean
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Build Python package
-      shell: bash -l {0}
-      env:
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        pip install build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-
-    - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash -l {0}
-      env:
-        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
-      run: |
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w

.github/actions/build-macos/action.yml

@@ -1,88 +0,0 @@
-name: 'Build and Test on macOS'
-description: 'Build and test MLX on macOS'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install dependencies
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
-      run: |
-        pip install --upgrade pip
-        pip install cmake setuptools nanobind==2.4.0
-        pip install -e . -v
-
-    - name: Generate package stubs
-      shell: bash -l {0}
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
-    - name: Install tests dependencies
-      shell: bash -l {0}
-      run: |
-        pip install numpy torch tensorflow unittest-xml-reporting
-
-    - name: Run Python tests
-      shell: bash -l {0}
-      env:
-        LOW_MEMORY: 1
-      run: |
-        DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-        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
-    
-    - name: Build example extension
-      shell: bash -l {0}
-      run: |
-        cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
-    - name: Build CPP only
-      shell: bash -l {0}
-      run: |
-        mkdir -p build
-        cd build
-        cmake ..
-        make -j $(sysctl -n hw.ncpu)
-    
-    - name: Run CPP tests
-      shell: bash -l {0}
-      env:
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
-    - name: Build small binary with JIT
-      shell: bash -l {0}
-      run: |
-        mkdir -p build
-        cd build
-        cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
-          -DBUILD_SHARED_LIBS=ON \
-          -DMLX_BUILD_CPU=OFF \
-          -DMLX_BUILD_SAFETENSORS=OFF \
-          -DMLX_BUILD_GGUF=OFF \
-          -DMLX_METAL_JIT=ON
-        make -j $(sysctl -n hw.ncpu)
-    
-    - name: Run Python tests with JIT
-      shell: bash -l {0}
-      env:
-        LOW_MEMORY: 1
-        DEVICE: gpu
-        METAL_DEVICE_WRAPPER_TYPE: 1
-        METAL_DEBUG_ERROR_MODE: 0
-      run: |
-        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
-        python -m xmlrunner discover \
-            -v python/tests \
-            -o test-results/gpu_jit

.github/actions/setup-linux/action.yml

@@ -1,86 +0,0 @@
-name: 'Setup Linux Environment'
-description: 'Install dependencies for Linux builds'
-
-inputs:
-  toolkit:
-    description: 'Which toolkit to install'
-    required: false
-    default: 'cpu'
-  python-version:
-    description: 'Version of python to set up'
-    required: false
-    default: '3.10'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Use ccache
-      uses: hendrikmuhs/ccache-action@v1.2
-      with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-py${{ inputs.python-version }}
-        max-size: 1GB
-
-    - name: Install common dependencies
-      shell: bash
-      run: |
-        sudo apt-get update
-        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip
-
-    - uses: actions/setup-python@v6
-      with:
-        python-version: ${{ inputs.python-version }}
-
-    - name: Setup Python venv
-      shell: bash
-      run: |
-        python -m venv .venv
-        source .venv/bin/activate
-        pip install setuptools cmake nanobind==2.4.0
-        echo PATH=$PATH >> $GITHUB_ENV
-        # Make cmake search .venv for nanobind
-        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
-
-    - name: Install MPI
-      shell: bash
-      run: sudo apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
-
-    - 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-toolkit-12-6",
-            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-toolkit-12-9",
-            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
-          }
-      run: |
-        # 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 \
-            libnccl2 libnccl-dev \
-            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
-        echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH
-
-    - name: CUDA packages and driver report
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      shell: bash
-      run: |
-        sudo apt-get install -y ubuntu-drivers-common dkms
-        echo "NVIDIA Driver Packages Available:"
-        sudo ubuntu-drivers list --gpgpu
-        echo "NVIDIA Driver Version:"
-        cat /proc/driver/nvidia/version || echo "nvidia driver not found"
-        echo "Installed NVIDIA and CUDA packages:"
-        dpkg -l | egrep "cuda|nvidia" -i
-        echo "DKMS Status:"
-        dkms status || echo "dkms not found"
-        echo "NVIDIA-SMI Status:"
-        nvidia-smi || echo "nvidia-smi not found"

.github/actions/setup-macos/action.yml

@@ -1,24 +0,0 @@
-name: 'Setup macOS Environment'
-description: 'Install dependencies for macOS builds'
-
-inputs:
-  python-version:
-    description: 'Python version to use'
-    required: false
-    default: '3.10'
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Homebrew packages
-      shell: sh
-      run: /opt/homebrew/bin/brew install openmpi
-    
-    - name: Verify MetalToolchain installed
-      shell: bash
-      run: xcodebuild -showComponent MetalToolchain
-
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}

.github/actions/test-linux/action.yml

@@ -1,69 +0,0 @@
-name: 'Run Linux tests'
-
-inputs:
-  has-gpu:
-    description: 'Run GPU tests'
-    required: false
-    default: false
-
-runs:
-  using: "composite"
-  steps:
-    - name: Run MPI tests
-      shell: bash
-      run: |
-        echo "::group::MPI tests"
-        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-        echo "::endgroup::"
-
-    - name: Run distributed tests
-      if: ${{ inputs.has-gpu == 'false' }}
-      shell: bash
-      run: |
-        echo "::group::Distributed tests"
-        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-        if grep -Fq '[WARN]' stderr.log ; then
-          grep -F '[WARN]' stderr.log
-          echo "Distributed ring test failed";
-          exit 1;
-        fi
-        echo "::endgroup::"
-
-    - name: Run Python tests - CPU
-      if: ${{ inputs.has-gpu == 'false' }}
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::Python tests - CPU"
-        python -m unittest discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run Python tests - GPU
-      if: ${{ inputs.has-gpu == 'true' }}
-      shell: bash
-      env:
-        DEVICE: gpu
-      run: |
-        echo "::group::Python tests - GPU"
-        python -m tests discover python/tests -v
-        echo "::endgroup::"
-
-    - name: Run CPP tests - CPU
-      shell: bash
-      env:
-        DEVICE: cpu
-      run: |
-        echo "::group::CPP tests - CPU"
-        ./build/tests/tests
-        echo "::endgroup::"
-
-    - name: Run CPP tests - GPU
-      if: ${{ inputs.has-gpu == 'true' }}
-      shell: bash
-      env:
-        DEVICE: gpu
-      run: |
-        echo "::group::CPP tests - GPU"
-        ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-        echo "::endgroup::"

.github/dependabot.yml

@@ -1,6 +0,0 @@
-version: 2
-updates:
-  - package-ecosystem: "github-actions"
-    directory: "/"
-    schedule:
-      interval: "weekly"

.github/scripts/setup+build-cpp-linux-fedora-container.sh

@@ -1,27 +0,0 @@
-#!/bin/bash
-set -ex
-
-# [Setup] Install dependencies inside the container.
-dnf update -y
-dnf install -y \
-  blas-devel \
-  lapack-devel \
-  openblas-devel \
-  make \
-  cmake \
-  clang \
-  git
-dnf clean all
-
-# [C++] CI Build Sanity Check: Verifies code compilation, not for release.
-export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-export DEBUG=1
-export CMAKE_C_COMPILER=/usr/bin/clang
-export CMAKE_CXX_COMPILER=/usr/bin/clang++
-
-mkdir -p build
-pushd build
-cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-make -j $(nproc)
-./tests/tests
-popd

.github/workflows/build_and_test.yml

@@ -1,108 +0,0 @@
-name: Build and Test
-
-on:
-  pull_request:
-  push:
-    branches:
-      - main
-      # For testing CI without starting a pull request:
-      - test/*
-
-permissions:
-  contents: read
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
-
-jobs:
-  check_lint:
-    name: Check Lint
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: pre-commit/action@v3.0.1
-
-  linux_build_and_test:
-    name: Linux (cpu, ${{ matrix.arch }})
-    needs: check_lint
-    strategy:
-      fail-fast: false
-      matrix:
-        arch: ['x86_64', 'aarch64']
-    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
-      - uses: ./.github/actions/test-linux
-
-  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"]
-    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
-
-  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_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

.github/workflows/documentation.yml

@@ -1,28 +0,0 @@
-name: Documentation
-
-on:
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build:
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/build-docs
-      
-  deploy:
-    needs: build
-    permissions:
-      pages: write
-      id-token: write
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4

.github/workflows/nightly.yml

@@ -1,96 +0,0 @@
-name: Nightly Build
-
-on:
-  schedule:
-    - cron: 33 6 * * 1-5
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-jobs:
-  build_linux_release:
-    strategy:
-      fail-fast: false
-      matrix:
-        python_version: ["3.10", "3.14"]
-    runs-on: ubuntu-22.04
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux-release
-        with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
-          arch: "x86_64"
-      - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v5
-        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@v5
-        with:
-          name: mlx-cpu
-          path: wheelhouse/mlx_cpu-*.whl
-          retention-days: 7
-
-  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
-
-  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 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'
-    runs-on: ubuntu-22-large
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: 'cuda-12.9'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-        with:
-          toolkit: 'cuda-12.9'
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
-          retention-days: 7

.github/workflows/pull_request.yml

@@ -0,0 +1,20 @@
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check_lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pre-commit black isort clang-format
+      - name: Run lint
+        run: |
+          pre-commit run --all-files

.github/workflows/release.yml

@@ -1,238 +0,0 @@
-name: PyPI Release
-
-on:
-  push:
-    tags:
-      - 'v*'
-  workflow_dispatch:
-    inputs:
-      dev_release:
-        description: "Do a dev release or regular release"
-        required: true
-        default: "false"
-
-permissions:
-  contents: read
-
-jobs:
-  setup:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Set publishing variables
-        run: echo "Publishing setup complete"
-
-  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:
-    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@v4
-
-  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: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          python-version: ${{ matrix.python_version }}
-      - 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@v5
-        with:
-          overwrite: true
-          name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
-          path: wheelhouse/mlx-*.whl
-      - name: Upload CPU artifacts
-        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v5
-        with:
-          overwrite: true
-          name: mlx-cpu-${{ matrix.arch }}
-          path: wheelhouse/mlx_cpu-*.whl
-  
-  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: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-macos
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools nanobind==2.4.0
-          pip install -e . -v
-      - name: Generate package stubs
-        shell: bash -l {0}
-        run: |
-          pip install typing_extensions
-          python setup.py generate_stubs
-      - 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: Upload MLX artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          overwrite: true
-          name: mac-wheels-${{ matrix.python-version }}
-          path: dist/mlx-*.whl
-      - name: Upload Metal artifacts
-        if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v5
-        with:
-          overwrite: true
-          name: mlx-metal
-          path: dist/mlx_metal-*.whl
-
-  build_cuda_release:
-    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
-    env:
-      PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-    steps:
-      - uses: actions/checkout@v6
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: 'cuda-12.9'
-      - name: Build Python package
-        uses: ./.github/actions/build-cuda-release
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v5
-        with:
-          overwrite: true
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
-
-  pypi-publish:
-    name: Upload release to PyPI
-    runs-on: ubuntu-latest
-    needs: [setup, build_linux_release, build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: pypi
-      url: https://pypi.org/p/mlx
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          pattern: linux-wheels-*
-          merge-multiple: true
-          path: dist
-      - uses: actions/download-artifact@v6
-        with:
-          pattern: mac-wheels-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        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: [setup, build_cuda_release]
-    permissions:
-      id-token: write
-    environment:
-      name: pypi
-      url: https://pypi.org/p/mlx-cuda
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          name: mlx-cuda
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        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: [setup, build_linux_release]
-    permissions:
-      id-token: write
-    environment:
-      name: pypi
-      url: https://pypi.org/p/mlx-cpu
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          pattern: mlx-cpu-*
-          merge-multiple: true
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        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: [setup, build_mac_release]
-    permissions:
-      id-token: write
-    environment:
-      name: pypi
-      url: https://pypi.org/p/mlx-metal
-    steps:
-      - uses: actions/download-artifact@v6
-        with:
-          name: mlx-metal
-          path: dist
-      - name: Display structure of downloaded files
-        run: ls -R dist
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          repository-url: https://upload.pypi.org/legacy/

.gitignore

@@ -36,7 +36,6 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
-uv.lock
 
 # vim
 *.swp
@@ -77,9 +76,6 @@ build/
 *.out
 *.app
 
-# Debug symbols
-*.pdb
-
 # VSCode 
 .vscode/
 .DS_Store

.pre-commit-config.yaml

@@ -1,22 +1,15 @@
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v6.0.0
-    hooks:
-    -   id: check-yaml
-    # -   id: end-of-file-fixer
-    # -   id: trailing-whitespace
 -   repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v19.1.7
+    rev: v18.1.8
     hooks:
     -   id: clang-format
 # Using this mirror lets us use mypyc-compiled black, which is about 2x faster
 -   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.1.0
+    rev: 24.8.0
     hooks:
     -   id: black
-    
 -   repo: https://github.com/pycqa/isort
-    rev: 6.0.0
+    rev: 5.13.2
     hooks:
     -   id: isort
         args:

ACKNOWLEDGMENTS.md

@@ -7,7 +7,7 @@ with a short description of your contribution(s) below. For example:
 
 MLX was developed with contributions from the following individuals:
 
-- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`. Added `orthogonal` initializer.
+- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`.
 - Juarez Bochi: Fixed bug in cross attention.
 - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example.
 - Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`.
@@ -19,17 +19,11 @@ MLX was developed with contributions from the following individuals:
 - Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
 - Paul Paczuski: Improved stability of BCE loss calculation
 - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops.
-- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.
 
 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
 </a>
 
-# Organizations
-
-MLX has received contributions from the following companies:
-- NVIDIA Corporation & Affiliates
-
 # Third-Party Software
 
 MLX leverages several third-party software, listed here together with

CMakeLists.txt

@@ -1,24 +1,6 @@
-cmake_minimum_required(VERSION 3.25)
+cmake_minimum_required(VERSION 3.24)
 
-if(NOT MLX_VERSION)
-  file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
-  string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}")
-  set(_major ${CMAKE_MATCH_1})
-  string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}")
-  set(_minor ${CMAKE_MATCH_1})
-  string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}")
-  set(_patch ${CMAKE_MATCH_1})
-  set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}")
-  set(MLX_VERSION ${MLX_PROJECT_VERSION})
-else()
-  string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION
-                       ${MLX_VERSION})
-endif()
-
-project(
-  mlx
-  LANGUAGES C CXX
-  VERSION ${MLX_PROJECT_VERSION})
+project(mlx LANGUAGES C CXX)
 
 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
@@ -26,7 +8,6 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
-set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # ----------------------------- Configuration -----------------------------
 option(MLX_BUILD_TESTS "Build tests for mlx" ON)
@@ -35,18 +16,19 @@ option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
 option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
 option(MLX_BUILD_METAL "Build metal backend" ON)
 option(MLX_BUILD_CPU "Build cpu backend" ON)
-option(MLX_BUILD_CUDA "Build cuda backend" OFF)
 option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
-option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 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)
+
+if(NOT MLX_VERSION)
+  set(MLX_VERSION 0.21.0)
+endif()
 
 # --------------------- Processor tests -------------------------
+
 message(
   STATUS
     "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}"
@@ -67,18 +49,10 @@ if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
       message(WARNING "Building for x86_64 arch is not officially supported.")
     endif()
   endif()
+
 else()
   set(MLX_BUILD_METAL OFF)
-endif()
-
-if(MLX_USE_CCACHE)
-  find_program(CCACHE_PROGRAM ccache)
-  if(CCACHE_PROGRAM)
-    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
-    set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-    set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
-  endif()
+  message(WARNING "MLX is prioritised for Apple silicon systems using macOS.")
 endif()
 
 # ----------------------------- Lib -----------------------------
@@ -89,26 +63,18 @@ cmake_policy(SET CMP0135 NEW)
 
 add_library(mlx)
 
-# Supress warnings: note: parameter passing for argument of type
-# ‘std::pair<float, float>’ when C++17 is enabled changed to match C++14 in GCC
-# 10.1
-target_compile_options(mlx PRIVATE -Wno-psabi)
-
-if(MLX_BUILD_CUDA)
-  enable_language(CUDA)
+if(MLX_BUILD_METAL)
+  set(METAL_LIB "-framework Metal")
+  set(FOUNDATION_LIB "-framework Foundation")
+  set(QUARTZ_LIB "-framework QuartzCore")
 endif()
 
-if(MLX_BUILD_METAL)
-  find_library(METAL_LIB Metal)
-  find_library(FOUNDATION_LIB Foundation)
-  find_library(QUARTZ_LIB QuartzCore)
-  if(METAL_LIB)
-    message(STATUS "Metal found ${METAL_LIB}")
-  else()
-    message(
-      FATAL_ERROR
-        "Metal not found. Set MLX_BUILD_METAL=OFF to build without GPU")
-  endif()
+if(MLX_BUILD_METAL AND NOT METAL_LIB)
+  message(STATUS "Metal not found. Unable to build GPU")
+  set(MLX_BUILD_METAL OFF)
+  set(MLX_METAL_DEBUG OFF)
+elseif(MLX_BUILD_METAL)
+  message(STATUS "Building METAL sources")
 
   if(MLX_METAL_DEBUG)
     add_compile_definitions(MLX_METAL_DEBUG)
@@ -117,8 +83,7 @@ if(MLX_BUILD_METAL)
   # Throw an error if xcrun not found
   execute_process(
     COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
-    OUTPUT_VARIABLE MACOS_SDK_VERSION
-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
+    OUTPUT_VARIABLE MACOS_SDK_VERSION COMMAND_ERROR_IS_FATAL ANY)
 
   if(${MACOS_SDK_VERSION} LESS 14.0)
     message(
@@ -128,12 +93,10 @@ if(MLX_BUILD_METAL)
   message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}")
 
   set(METAL_CPP_URL
-      https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip)
+      https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18-beta.zip
+  )
 
   if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
-    if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0)
-      message(FATAL_ERROR "MLX requires macOS >= 14.0")
-    endif()
     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
   endif()
   execute_process(
@@ -142,6 +105,7 @@ if(MLX_BUILD_METAL)
       "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
     OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+
   FetchContent_MakeAvailable(metal_cpp)
   target_include_directories(
     mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
@@ -149,62 +113,16 @@ if(MLX_BUILD_METAL)
   target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
 endif()
 
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
-  # With newer clang/gcc versions following libs are implicitly linked, but when
-  # building on old distributions they need to be explicitly listed.
-  target_link_libraries(mlx PRIVATE dl pthread)
-endif()
-
-if(WIN32)
-  if(MSVC)
-    # GGUF does not build with MSVC.
-    set(MLX_BUILD_GGUF OFF)
-    # There is no prebuilt OpenBLAS distribution for MSVC.
-    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
-  endif()
-  # Windows implementation of dlfcn.h APIs.
-  FetchContent_Declare(
-    dlfcn-win32
-    GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git
-    GIT_TAG v1.4.1
-    EXCLUDE_FROM_ALL)
-  block()
-  set(BUILD_SHARED_LIBS OFF)
-  FetchContent_MakeAvailable(dlfcn-win32)
-  endblock()
-  target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src")
-  target_link_libraries(mlx PRIVATE dl)
-endif()
-
 if(MLX_BUILD_CPU)
   find_library(ACCELERATE_LIBRARY Accelerate)
   if(ACCELERATE_LIBRARY)
     message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
     set(MLX_BUILD_ACCELERATE ON)
-  else()
-    message(STATUS "Accelerate not found, using default backend.")
-    set(MLX_BUILD_ACCELERATE OFF)
-  endif()
-
-  if(MLX_BUILD_ACCELERATE)
     target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
-    add_compile_definitions(MLX_USE_ACCELERATE)
     add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
-    # Download and build OpenBLAS from source code.
-    FetchContent_Declare(
-      openblas
-      GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git
-      GIT_TAG v0.3.28
-      EXCLUDE_FROM_ALL)
-    set(BUILD_STATIC_LIBS ON) # link statically
-    set(NOFORTRAN ON) # msvc has no fortran compiler
-    FetchContent_MakeAvailable(openblas)
-    target_link_libraries(mlx PRIVATE openblas)
-    target_include_directories(
-      mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include"
-                  "${CMAKE_BINARY_DIR}/generated" "${CMAKE_BINARY_DIR}")
   else()
+    message(STATUS "Accelerate or arm neon not found, using default backend.")
+    set(MLX_BUILD_ACCELERATE OFF)
     if(${CMAKE_HOST_APPLE})
       # The blas shipped in macOS SDK is not supported, search homebrew for
       # openblas instead.
@@ -222,7 +140,7 @@ if(MLX_BUILD_CPU)
     message(STATUS "Lapack lib " ${LAPACK_LIBRARIES})
     message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS})
     target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS})
-    target_link_libraries(mlx PRIVATE ${LAPACK_LIBRARIES})
+    target_link_libraries(mlx PUBLIC ${LAPACK_LIBRARIES})
     # List blas after lapack otherwise we may accidentally incldue an old
     # version of lapack.h from the include dirs of blas.
     find_package(BLAS REQUIRED)
@@ -235,19 +153,36 @@ if(MLX_BUILD_CPU)
     message(STATUS "Blas lib " ${BLAS_LIBRARIES})
     message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS})
     target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS})
-    target_link_libraries(mlx PRIVATE ${BLAS_LIBRARIES})
+    target_link_libraries(mlx PUBLIC ${BLAS_LIBRARIES})
+
+    if(WIN32)
+      find_package(dlfcn-win32 REQUIRED)
+      message(STATUS "dlfcn-win32 lib " ${dlfcn-win32_LIBRARIES})
+      message(STATUS "dlfcn-win32 include " ${dlfcn-win32_INCLUDE_DIRS})
+      target_link_libraries(mlx PUBLIC ${dlfcn-win32_LIBRARIES})
+    endif()
   endif()
 else()
   set(MLX_BUILD_ACCELERATE OFF)
 endif()
 
-message(STATUS "Downloading json")
-FetchContent_Declare(
-  json
-  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
-FetchContent_MakeAvailable(json)
-target_include_directories(
-  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
+find_package(MPI)
+if(MPI_FOUND)
+  execute_process(
+    COMMAND zsh "-c" "mpirun --version"
+    OUTPUT_VARIABLE MPI_VERSION
+    ERROR_QUIET)
+  if(${MPI_VERSION} MATCHES ".*Open MPI.*")
+    target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
+  elseif(MPI_VERSION STREQUAL "")
+    set(MPI_FOUND FALSE)
+    message(
+      WARNING "MPI found but mpirun is not available. Building without MPI.")
+  else()
+    set(MPI_FOUND FALSE)
+    message(WARNING "MPI which is not OpenMPI found. Building without MPI.")
+  endif()
+endif()
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
 
@@ -255,19 +190,12 @@ 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(
+FetchContent_Declare(
   fmt
   GIT_REPOSITORY https://github.com/fmtlib/fmt.git
   GIT_TAG 10.2.1
   EXCLUDE_FROM_ALL)
-  FetchContent_MakeAvailable(fmt)
-endif()
+FetchContent_MakeAvailable(fmt)
 target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>)
 
 if(MLX_BUILD_PYTHON_BINDINGS)
@@ -279,7 +207,8 @@ if(MLX_BUILD_PYTHON_BINDINGS)
   execute_process(
     COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
     OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE nanobind_ROOT)
+    OUTPUT_VARIABLE NB_DIR)
+  list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
   find_package(nanobind CONFIG REQUIRED)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()

CONTRIBUTING.md

@@ -17,11 +17,11 @@ possible.
  
    You can also run the formatters manually as follows:
  
-   ```shell
+     ```
      clang-format -i file.cpp
      ```
  
-   ```shell
+     ```
      black file.py
      ```
  

MANIFEST.in

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

README.md

@@ -11,28 +11,28 @@ brought to you by Apple machine learning research.
 
 Some key features of MLX include:
 
-- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+ - **Familiar APIs**: MLX has a Python API that closely follows NumPy.  MLX
    also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
    [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
    the Python API.  MLX has higher-level packages like `mlx.nn` and
    `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
    more complex models.
 
-- **Composable function transformations**: MLX supports composable function
+ - **Composable function transformations**: MLX supports composable function
    transformations for automatic differentiation, automatic vectorization,
    and computation graph optimization.
 
-- **Lazy computation**: Computations in MLX are lazy. Arrays are only
+ - **Lazy computation**: Computations in MLX are lazy. Arrays are only
    materialized when needed.
 
-- **Dynamic graph construction**: Computation graphs in MLX are constructed
+ - **Dynamic graph construction**: Computation graphs in MLX are constructed
    dynamically. Changing the shapes of function arguments does not trigger
    slow compilations, and debugging is simple and intuitive.
 
-- **Multi-device**: Operations can run on any of the supported devices
+ - **Multi-device**: Operations can run on any of the supported devices
    (currently the CPU and the GPU).
 
-- **Unified memory**: A notable difference from MLX and other frameworks
+ - **Unified memory**: A notable difference from MLX and other frameworks
    is the *unified memory model*. Arrays in MLX live in shared memory.
    Operations on MLX arrays can be performed on any of the supported
    device types without transferring data.
@@ -68,23 +68,18 @@ in the documentation.
 
 ## Installation
 
-MLX is available on [PyPI](https://pypi.org/project/mlx/). To install MLX on
-macOS, run:
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install the Python API, run:
 
-```bash
+**With `pip`**:
+
+```
 pip install mlx
 ```
 
-To install the CUDA backend on Linux, run:
+**With `conda`**:
 
-```bash
-pip install mlx[cuda]
 ```
-
-To install a CPU-only Linux package, run:
-
-```bash
-pip install mlx[cpu]
+conda install -c conda-forge mlx
 ```
 
 Checkout the
@@ -110,7 +105,7 @@ Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
 MLX useful in your research and wish to cite it, please use the following
 BibTex entry:
 
-```text
+```
 @software{mlx2023,
   author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
   title = {{MLX}: Efficient and flexible machine learning on Apple silicon},

benchmarks/cpp/autograd.cpp

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

benchmarks/cpp/compare_devices.cpp

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

benchmarks/cpp/irregular_strides.cpp

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

benchmarks/cpp/single_ops.cpp

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

benchmarks/python/blas/bench_gemm.py

@@ -142,7 +142,9 @@ def bench_shape(B, M, N, K, np_dtype, transpose="nn"):
     t_b = (0, 1, 2) if transpose[1] == "n" else (0, 2, 1)
 
     c_mlx = a_mx.transpose(t_a) @ b_mx.transpose(t_b)
-    c_npy = a_np.transpose(t_a).astype(np_dtype) @ b_np.transpose(t_b).astype(np_dtype)
+    c_npy = a_np.transpose(t_a).astype(np.float32) @ b_np.transpose(t_b).astype(
+        np.float32
+    )
 
     atol = 1e-5 if np_dtype == np.float32 else 1e-4
 
@@ -161,7 +163,7 @@ def get_gflop_count(B, M, N, K):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run gemm benchmarks")
 
-    dtypes = ("float32", "float16", "complex64")
+    dtypes = ("float32", "float16")
     transposes = ("nn", "nt", "tn")
     shapes = (
         (16, 234, 768, 3072),
@@ -185,7 +187,7 @@ if __name__ == "__main__":
                 diff = gflops_mx / gflops_pt - 1.0
 
                 print(
-                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100.0 * diff:+5.2f}%"
+                    f"{B:3d}, {M:4d}, {N:4d}, {K:4d}, {dtype}, {transpose}, {gflops_pt:05.3f}, {gflops_mx:05.3f}, {100. * diff:+5.2f}%"
                 )
                 if gflops_pt >= 2.0 * gflops_mx:
                     print("ATTENTION ^^^^^^^")

benchmarks/python/blas/bench_gemv.py

@@ -1,5 +1,6 @@
 # Copyright © 2023 Apple Inc.
 
+import argparse
 import os
 import subprocess
 import time
@@ -195,7 +196,7 @@ def bench_with_out_len(ax, out_vec_len, in_vector_lens, dtype, transpose):
 
 
 for transpose in (False, True):
-    for dtype in ("float32", "float16", "complex64"):
+    for dtype in ("float32", "float16"):
         fig, axs = plt.subplots(
             len(in_vec_sizes), 2, figsize=(8.5, 11), layout="constrained"
         )
@@ -214,7 +215,7 @@ for transpose in (False, True):
         fig.suptitle(f"{device_name}: {dtype} {op_name}")
         fig.savefig(
             os.path.join(
-                results_dir, f"{device_name.replace(' ', '_')}_{dtype}_{op_name}.pdf"
+                results_dir, f'{device_name.replace(" ", "_")}_{dtype}_{op_name}.pdf'
             )
         )
         plt.close(fig)

benchmarks/python/comparative/bench_torch.py

@@ -5,7 +5,6 @@ import os
 import time
 
 import torch
-import torch.cuda
 import torch.mps
 
 
@@ -45,10 +44,8 @@ def bench(f, *args):
 
 
 def sync_if_needed(x):
-    if x.device == torch.device("mps"):
+    if x.device != torch.device("cpu"):
         torch.mps.synchronize()
-    elif x.device == torch.device("cuda"):
-        torch.cuda.synchronize()
 
 
 @torch.no_grad()
@@ -102,14 +99,6 @@ def reduction(op, axis, x):
     sync_if_needed(x)
 
 
-@torch.no_grad()
-def sum_and_add(axis, x, y):
-    z = x.sum(axis=axis, keepdims=True)
-    for i in range(50):
-        z = (z + y).sum(axis=axis, keepdims=True)
-    sync_if_needed(x)
-
-
 @torch.no_grad()
 def softmax(axis, x):
     ys = []
@@ -351,11 +340,7 @@ if __name__ == "__main__":
         args.axis.pop(0)
 
     torch.set_num_threads(1)
-    device = "mps"
-    if torch.cuda.is_available():
-        device = "cuda"
-    if args.cpu:
-        device = "cpu"
+    device = "cpu" if args.cpu else "mps"
 
     types = args.dtype
     if not types:
@@ -475,8 +460,5 @@ if __name__ == "__main__":
     elif args.benchmark == "selu":
         print(bench(selu, x))
 
-    elif args.benchmark == "sum_and_add":
-        print(bench(sum_and_add, axis, *xs))
-
     else:
         raise ValueError(f"Unknown benchmark `{args.benchmark}`.")

benchmarks/python/conv_unaligned_bench.py

@@ -1,107 +0,0 @@
-import math
-import time
-
-import mlx.core as mx
-import numpy as np
-import torch
-
-N_warmup = 10
-N_iter_bench = 100
-N_iter_func = 5
-
-
-def bench(f, a, b):
-    for i in range(N_warmup):
-        f(a, b)
-    torch.mps.synchronize()
-
-    s = time.perf_counter_ns()
-    for i in range(N_iter_bench):
-        f(a, b)
-    e = time.perf_counter_ns()
-    return (e - s) * 1e-9
-
-
-def make_mx_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
-    def mx_conv_2D(a, b):
-        ys = []
-        for i in range(N_iter_func):
-            y = mx.conv2d(a, b, stride=strides, padding=padding, groups=groups)
-            ys.append(y)
-        mx.eval(ys)
-        return ys
-
-    return mx_conv_2D
-
-
-def make_pt_conv_2D(strides=(1, 1), padding=(0, 0), groups=1):
-    @torch.no_grad()
-    def pt_conv_2D(a, b):
-        ys = []
-        for i in range(N_iter_func):
-            y = torch.conv2d(a, b, stride=strides, padding=padding, groups=groups)
-            ys.append(y)
-        torch.mps.synchronize()
-        return ys
-
-    return pt_conv_2D
-
-
-def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype):
-    scale = 1.0 / math.sqrt(kH * kH * C)
-    a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype)
-    b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype(
-        np_dtype
-    )
-
-    a_mx = mx.array(a_np)
-    b_mx = mx.array(b_np)
-
-    a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("mps")
-    b_pt = torch.from_numpy(b_np.transpose((0, 3, 1, 2))).to("mps")
-
-    torch.mps.synchronize()
-
-    f_mx = make_mx_conv_2D(strides, padding, groups)
-    f_pt = make_pt_conv_2D(strides, padding, groups)
-
-    time_torch = bench(f_pt, a_pt, b_pt)
-    time_mlx = bench(f_mx, a_mx, b_mx)
-
-    out_mx = mx.conv2d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
-    out_pt = torch.conv2d(
-        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
-    )
-    out_pt = torch.permute(out_pt, (0, 2, 3, 1))
-    out_pt = out_pt.numpy(force=True)
-
-    atol = 2e-5 if np_dtype == np.float32 else 1e-4
-
-    if not np.allclose(out_pt, out_mx, atol=atol):
-        print(
-            f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
-        )
-
-    return time_mlx, time_torch
-
-
-if __name__ == "__main__":
-    dtype = "float32"
-    shapes = (
-        (4, 32, 32, 21, 3, 3, 128),
-        (4, 32, 32, 21, 3, 3, 37),
-        (4, 32, 32, 370, 3, 3, 370),
-        (4, 32, 32, 370, 7, 7, 128),
-        (2, 320, 640, 21, 7, 7, 21),
-    )
-    for N, H, W, C, kh, kw, O in shapes:
-        time_mlx, time_torch = bench_shape(
-            N, H, W, C, kh, kw, O, (1, 1), (0, 0), 1, dtype
-        )
-        diff = time_torch / time_mlx - 1.0
-
-        print(
-            f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kh:2d}, {kw:2d}, {C:3d}), {dtype}, {100. * diff:+5.2f}%"
-        )
-        if time_mlx >= 2.0 * time_torch:
-            print("ATTENTION ^^^^^^^")

benchmarks/python/gather_bench.py

@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+from time import time
 
 import mlx.core as mx
 import torch

benchmarks/python/gather_mm_bench.py

@@ -1,74 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate([x[i] @ w[j].T for i, j in enumerate(idx.tolist())], axis=0)
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_mm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_mm(x, w1.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_mm(x, w2.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = x @ w1.T
-        x = x @ w2.T
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_mm()

benchmarks/python/gather_qmm_bench.py

@@ -1,84 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import mlx.core as mx
-from time_utils import time_fn
-
-N = 1024
-D = 1024
-M = 1024
-E = 32
-I = 4
-
-
-def gather_sort(x, indices):
-    N, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
-def gather_mm_simulate(x, w, indices):
-    x, idx, inv_order = gather_sort(x, indices)
-    for i in range(2):
-        y = mx.concatenate(
-            [
-                mx.quantized_matmul(x[i], w[0][j], w[1][j], w[2][j], transpose=True)
-                for i, j in enumerate(idx.tolist())
-            ],
-            axis=0,
-        )
-        x = y[:, None]
-    x = scatter_unsort(x, inv_order, indices.shape)
-    return x
-
-
-def time_gather_qmm():
-    x = mx.random.normal((N, 1, 1, D)) / 1024**0.5
-    w1 = mx.random.normal((E, M, D)) / 1024**0.5
-    w2 = mx.random.normal((E, D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32)
-    sorted_indices = mx.sort(indices.flatten()).reshape(N, I)
-    mx.eval(x, w1, w2, indices, sorted_indices)
-
-    def gather_mm(x, w1, w2, indices, sort):
-        idx = indices
-        inv_order = None
-        if sort:
-            x, idx, inv_order = gather_sort(x, indices)
-        x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=sort)
-        if sort:
-            x = scatter_unsort(x, inv_order, indices.shape)
-        return x
-
-    time_fn(gather_mm, x, w1, w2, indices, False)
-    time_fn(gather_mm, x, w1, w2, sorted_indices, False)
-    time_fn(gather_mm, x, w1, w2, indices, True)
-
-    x = mx.random.normal((N * I, D)) / 1024**0.5
-    w1 = mx.random.normal((M, D)) / 1024**0.5
-    w2 = mx.random.normal((D, M)) / 1024**0.5
-    w1 = mx.quantize(w1)
-    w2 = mx.quantize(w2)
-    mx.eval(x, w1, w2)
-
-    def equivalent_matmul(x, w1, w2):
-        x = mx.quantized_matmul(x, *w1, transpose=True)
-        x = mx.quantized_matmul(x, *w2, transpose=True)
-        return x
-
-    time_fn(equivalent_matmul, x, w1, w2)
-
-
-if __name__ == "__main__":
-    time_gather_qmm()

benchmarks/python/layer_norm_bench.py

@@ -1,7 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.
 
-from functools import partial
-
 import mlx.core as mx
 import mlx.nn as nn
 from time_utils import time_fn
@@ -12,71 +10,32 @@ def layer_norm(x, w, b, eps):
     x = x.astype(mx.float32)
     mu = mx.mean(x, -1, keepdims=True)
     v = mx.var(x, -1, keepdims=True)
-    y = (x - mu) * mx.rsqrt(v + eps)
-    if w is not None:
-        y = y * w
-    if b is not None:
-        y = y + b
-    return y
+    return (x - mu) * mx.rsqrt(v + eps) * w + b
 
 
-def time_layer_norm(N, dt):
-    L = 1024
+def time_layer_norm():
     f1 = lambda x, w, b, y: (layer_norm(x, w, b, 1e-5) * y).sum()
     f2 = lambda x, w, b, y: (mx.fast.layer_norm(x, w, b, 1e-5) * y).sum()
     g1 = mx.grad(f1, argnums=(0, 1, 2))
     g2 = mx.grad(f2, argnums=(0, 1, 2))
 
-    x = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    w = mx.random.uniform(shape=(N,)).astype(dt)
-    b = mx.random.uniform(shape=(N,)).astype(dt)
-    y = mx.random.uniform(shape=(8, L, N)).astype(dt)
+    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
+    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    b = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
     mx.eval(x, w, b, y)
 
-    def layer_norm_loop(f, x, w, b):
-        for _ in range(32):
-            x = f(x, w, b)
-        return x
-
-    time_fn(layer_norm_loop, partial(layer_norm, eps=1e-5), x, w, b)
-    time_fn(layer_norm_loop, partial(mx.fast.layer_norm, eps=1e-5), x, w, b)
-
-    def layer_norm_grad_loop(g, x, w, b):
+    def layer_norm_loop(g, x, w, b):
         gx, gw, gb = x, w, b
         for _ in range(32):
             gx, gw, gb = g(gx, gw, gb, y)
         return gx, gw, gb
 
-    time_fn(layer_norm_grad_loop, g1, x, w, b)
-    time_fn(layer_norm_grad_loop, g2, x, w, b)
-    time_fn(layer_norm_grad_loop, mx.compile(g1), x, w, b)
-    time_fn(layer_norm_grad_loop, mx.compile(g2), x, w, b)
-
-    f1 = lambda x, y: (layer_norm(x, None, None, 1e-5) * y).sum()
-    f2 = lambda x, y: (mx.fast.layer_norm(x, None, None, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0,))
-    g2 = mx.grad(f2, argnums=(0,))
-
-    x = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    w = mx.random.uniform(shape=(N,)).astype(dt)
-    b = mx.random.uniform(shape=(N,)).astype(dt)
-    y = mx.random.uniform(shape=(8, L, N)).astype(dt)
-    mx.eval(x, w, b, y)
-
-    def layer_norm_grad_x_loop(g, x):
-        gx = x
-        for _ in range(32):
-            gx = g(gx, y)
-        return gx
-
-    time_fn(layer_norm_grad_x_loop, g1, x)
-    time_fn(layer_norm_grad_x_loop, g2, x)
-    time_fn(layer_norm_grad_x_loop, mx.compile(g1), x)
-    time_fn(layer_norm_grad_x_loop, mx.compile(g2), x)
+    time_fn(layer_norm_loop, g1, x, w, b)
+    time_fn(layer_norm_loop, g2, x, w, b)
+    time_fn(layer_norm_loop, mx.compile(g1), x, w, b)
+    time_fn(layer_norm_loop, mx.compile(g2), x, w, b)
 
 
 if __name__ == "__main__":
-    for dt in [mx.float32, mx.float16, mx.bfloat16]:
-        for n in [1024, 2048, 4096, 8192, 8192 + 1024]:
-            print(dt, n)
-            time_layer_norm(n, dt)
+    time_layer_norm()

benchmarks/python/masked_scatter.py

@@ -1,212 +0,0 @@
-import math
-import os
-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)
-
-DEVICE_NAME = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
-DEVICE_NAME = DEVICE_NAME.decode("utf-8").strip("\n")
-
-TORCH_DEVICE = torch.device(
-    "mps"
-    if torch.backends.mps.is_available()
-    else ("cuda" if torch.cuda.is_available() else "cpu")
-)
-
-
-N_WARMUP = 5
-N_ITER_BENCH = 50
-N_ITER_FUNC = 20
-
-VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)]
-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}.pdf",
-        )
-        fig.savefig(output_path)
-        plt.close(fig)
-
-
-if __name__ == "__main__":
-    main()

benchmarks/python/rms_norm_bench.py

@@ -9,10 +9,7 @@ def rms_norm(x, w, eps):
     ot = x.dtype
     x = x.astype(mx.float32)
     n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps)
-    y = (x * n).astype(ot)
-    if w is not None:
-        y = y * w
-    return y
+    return (x * n).astype(ot) * w
 
 
 def time_rms_norm():
@@ -37,27 +34,6 @@ def time_rms_norm():
     time_fn(rms_norm_loop, mx.compile(g1), x, w)
     time_fn(rms_norm_loop, mx.compile(g2), x, w)
 
-    f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum()
-    f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum()
-    g1 = mx.grad(f1, argnums=(0,))
-    g2 = mx.grad(f2, argnums=(0,))
-
-    x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    w = mx.random.uniform(shape=(4096,)).astype(mx.float16)
-    y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16)
-    mx.eval(x, w, y)
-
-    def rms_norm_loop(g, x):
-        gx = x
-        for _ in range(32):
-            gx = g(gx, y)
-        return gx
-
-    time_fn(rms_norm_loop, g1, x)
-    time_fn(rms_norm_loop, g2, x)
-    time_fn(rms_norm_loop, mx.compile(g1), x)
-    time_fn(rms_norm_loop, mx.compile(g2), x)
-
 
 if __name__ == "__main__":
     time_rms_norm()

benchmarks/python/sdpa_bench.py

@@ -28,34 +28,11 @@ def bench(f, *args):
     return (e - s) * 1e-9
 
 
-def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype):
-    np_dtype = getattr(np, dtype)
-
-    shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D)
-    shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D)
-
-    scale = 1.0 / math.sqrt(D)
-
-    q_np = np.random.normal(0.0, 1.0, shape_q).astype(np_dtype)
-    k_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-    v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype)
-
-    q_mx = mx.array(q_np)
-    k_mx = mx.array(k_np)
-    v_mx = mx.array(v_np)
-
-    if mask is not None:
-        if mask == "additive":
-            mask_np = np.random.normal(0.0, 1.0, (B, qH, qL, kL)).astype(np_dtype)
-            mask = mx.array(mask_np)
-        elif mask == "bool":
-            mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5
-            mask = mx.array(mask_np)
-
-    return q_mx, k_mx, v_mx, scale, mask
+def mlx_sdpa_fused_inner(q, k, v, scale):
+    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None)
 
 
-def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
+def mlx_sdpa_unfused_inner(q, k, v, scale, f32softmax=False):
     q_dtype = q.dtype
     q = q * mx.array(scale, q_dtype)
     n_q_heads = q.shape[-3]
@@ -64,7 +41,6 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
 
     B = q.shape[0]
     L = q.shape[2]
-    kL = k.shape[2]
 
     if n_repeats > 1:
         q = mx.reshape(q, [B, n_kv_heads, n_repeats, L, -1])
@@ -72,27 +48,10 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
         v = mx.expand_dims(v, 2)
 
     scores = q @ mx.swapaxes(k, -1, -2)
-
-    if mask is not None:
-
-        if mask == "causal":
-            q_offset = max(0, kL - L)
-            q_indices = mx.arange(q_offset, q_offset + L)
-            k_indices = mx.arange(kL)
-            mask = q_indices[:, None] >= k_indices[None]
-
-        if n_repeats > 1 and mask.ndim >= 3:
-            if mask.shape[-3] == 1:
-                mask = mx.expand_dims(mask, -3)
+    if f32softmax:
+        scores = mx.softmax(scores.astype(mx.float32), axis=-1).astype(q_dtype)
     else:
-                mask = mx.unflatten(mask, -3, (n_kv_heads, n_repeats))
-
-        if mask.dtype == mx.bool_:
-            scores = mx.where(mask, scores, -np.float32(np.inf))
-        else:
-            scores += mask
-
-    scores = mx.softmax(scores, axis=-1, precise=True)
+        scores = mx.softmax(scores, axis=-1)
 
     out = scores @ v
     if n_repeats > 1:
@@ -101,55 +60,74 @@ def mlx_ref_attn(q, k, v, scale=1.0, mask=None):
     return out
 
 
-def mlx_fused_attn(q, k, v, scale, mask):
-    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention(f, q, k, v, scale, mask=None, transpose=False):
-    if transpose:
-        q_t = mx.transpose(q, (0, 2, 1, 3))
-        k_t = mx.transpose(k, (0, 2, 1, 3))
-        v_t = mx.transpose(v, (0, 2, 1, 3))
-        o_t = f(q_t, k_t, v_t, scale=scale, mask=mask)
-        return mx.transpose(o_t, (0, 2, 1, 3))
-    else:
-        return f(q, k, v, scale=scale, mask=mask)
-
-
-def do_attention_bench(f, q, k, v, scale, mask=None, transpose=False):
+def mlx_spda_unfused(q, k, v, scale, transpose):
     q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
 
     for i in range(N_iter_func):
-        q_out = do_attention(f, q_out, k, v, scale, mask=mask, transpose=transpose)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_unfused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
 
     mx.eval(q_out)
     return q_out
 
 
-def bench_shape(
-    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, dtype, transpose=True, mask_in=None
-):
-    q_mx, k_mx, v_mx, scale, mask = prepare_inputs(
-        B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, mask_in, transpose, dtype
+def mlx_spda_fused(q, k, v, scale, transpose):
+    q_out = q
+    if transpose:
+        k = mx.transpose(k, (0, 2, 1, 3))
+        v = mx.transpose(v, (0, 2, 1, 3))
+
+    for i in range(N_iter_func):
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+        q_out = mlx_sdpa_fused_inner(q_out, k, v, scale)
+        if transpose:
+            q_out = mx.transpose(q_out, (0, 2, 1, 3))
+
+    mx.eval(q_out)
+    return q_out
+
+
+def bench_shape(B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose=True):
+    shape_q = (
+        (B, qsl, n_q_heads, head_dim) if transpose else (B, n_q_heads, qsl, head_dim)
+    )
+    shape_kv = (
+        (B, ksl, n_kv_heads, head_dim) if transpose else (B, n_kv_heads, ksl, head_dim)
     )
 
-    time_mlx_unfused = bench(
-        do_attention_bench, mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
-    time_mlx_fused = bench(
-        do_attention_bench, mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    q_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_q).astype(np_dtype)
+    k_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
+    v_np = np.random.normal(0.0, 1.0 / math.sqrt(head_dim), shape_kv).astype(np_dtype)
 
-    o_mlx_fused = do_attention(mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, transpose)
-    o_mlx_unfused = do_attention(
-        mlx_fused_attn, q_mx, k_mx, v_mx, scale, mask, transpose
-    )
+    scale = math.sqrt(1.0 / head_dim)
 
-    atol = 1e-5 if dtype == "float32" else 2e-4
+    q_mx = mx.array(q_np)
+    k_mx = mx.array(k_np)
+    v_mx = mx.array(v_np)
 
-    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol, rtol=atol):
+    time_mlx_unfused = bench(mlx_spda_unfused, q_mx, k_mx, v_mx, scale, transpose)
+    time_mlx_fused = bench(mlx_spda_fused, q_mx, k_mx, v_mx, scale, transpose)
+
+    if transpose:
+        q_mx = mx.transpose(q_mx, (0, 2, 1, 3))
+        k_mx = mx.transpose(k_mx, (0, 2, 1, 3))
+        v_mx = mx.transpose(v_mx, (0, 2, 1, 3))
+
+    o_mlx_fused = mlx_sdpa_fused_inner(q_mx, k_mx, v_mx, scale)
+    o_mlx_unfused = mlx_sdpa_unfused_inner(q_mx, k_mx, v_mx, scale, f32softmax=True)
+
+    atol = 1e-5 if np_dtype == np.float32 else 1e-4
+
+    if not mx.allclose(o_mlx_fused, o_mlx_unfused, atol=atol):
         print(
-            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}, mask: {mask_in}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
+            f"Failed at (B: {B}, qsl: {qsl}, ksl: {ksl}, head_dim: {head_dim}, n_qh: {n_q_heads}, n_kvh: {n_kv_heads}) [tpose = {transpose}] with max(|a - b|) = {mx.max(mx.abs(o_mlx_unfused - o_mlx_fused)):3.2e}"
         )
 
     return time_mlx_fused, time_mlx_unfused
@@ -173,51 +151,39 @@ if __name__ == "__main__":
           (  1,   128,   128,       64,   32,    32),
           (  1,   256,   256,       64,   32,    32),
           (  1,   512,   512,       64,   32,    32),
-          (  1,  1024,  1024,       64,   32,     8),
-          (  1,  2048,  2048,       64,   32,     8),
-          (  1,  4096,  4096,       64,   32,     8),
+          (  1,  1024,  1024,       64,   32,    32),
+          (  1,  2048,  2048,       64,   32,    32),
+          (  1,  4096,  4096,       64,   32,    32),
     )
 
     shapes_80 = (
         # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
-          (  1,  1024,  1024,       80,   32,     8),
-          (  1,  2048,  2048,       80,   32,     8),
-          (  1,  4096,  4096,       80,   32,     8),
+          (  1,  1024,  1024,       80,   32,    32),
+          (  1,  2048,  2048,       80,   32,    32),
+          (  1,  4096,  4096,       80,   32,    32),
     )
 
     shapes_128 = (
         # (  B,   qsl,   ksl, head_dim, n_qh, n_kvh)
-          (  1,  1024,  1024,      128,   32,     8),
-          (  1,  2048,  2048,      128,   32,     8),
-          (  1,  4096,  4096,      128,   32,     8),
+          (  1,  1024,  1024,      128,   32,    32),
+          (  1,  2048,  2048,      128,   32,    32),
+          (  1,  4096,  4096,      128,   32,    32),
     )
     # fmt: on
 
     shapes = shapes_64 + shapes_80 + shapes_128
 
-    masks = [None, "bool", "causal"]
-
-    print(
-        "  B,   qsl,   ksl, hdim, n_qh, n_kvh, t,   dtype,     mask, t_unfs, t_fuse, diff%"
-    )
+    print("  B,   qsl,   ksl, hdim, n_qh, n_kvh, tpose,   dtype, t_unfs, t_fuse, diff%")
 
     for dtype in dtypes:
         for transpose in transposes:
             for B, qsl, ksl, head_dim, n_q_heads, n_kv_heads in shapes:
-                for mask_in in masks:
+                np_dtype = getattr(np, dtype)
                 time_mlx_fused, time_mlx_unfused = bench_shape(
-                        B,
-                        qsl,
-                        ksl,
-                        head_dim,
-                        n_q_heads,
-                        n_kv_heads,
-                        dtype,
-                        transpose,
-                        mask_in,
+                    B, qsl, ksl, head_dim, n_q_heads, n_kv_heads, np_dtype, transpose
                 )
                 diff = time_mlx_unfused / time_mlx_fused - 1.0
                 t_str = 1 if transpose else 0
                 print(
-                        f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:1d}, {dtype}, {str(mask_in):>8}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
+                    f"{B:3d}, {qsl:5d}, {ksl:5d}, {head_dim:4d}, {n_q_heads:4d}, {n_kv_heads:5d}, {t_str:5d}, {dtype}, {time_mlx_unfused: 2.3f}, {time_mlx_fused: 2.3f}, {100. * diff:+5.2f}%"
                 )

benchmarks/python/sdpa_vector_bench.py

@@ -1,95 +1,94 @@
-import argparse
-import math
-
 import mlx.core as mx
+import numpy as np
+from mlx.utils import tree_map
 from time_utils import time_fn
 
-L = 16384
+L = 32768
 H = 32
 H_k = H // 4
 D = 128
-V = 128
 dtype = mx.float16
-loops = 10
+bits = 8
+
+loops = 20
 
 
-def upproject(x, w):
-    if w is None:
-        return x
-    else:
-        return x @ w.T
-
-
-def attention(q, k, v, mask=None, w=None):
-    def _sdpa(q, k, v):
+def attention(q, k, v):
+    for _ in range(loops):
         B, Hq, L, D = q.shape
         _, Hk, S, _ = k.shape
-        _, _, _, V = v.shape
         q = q.reshape(B, Hk, Hq // Hk, L, D)
-        k = k[:, :, None, :, :]
-        v = v[:, :, None, :, :]
-        s = q @ k.transpose(0, 1, 2, 4, 3)
-        if mask is not None:
-            m = mx.broadcast_to(mask, (B, Hq, L, S)).reshape(B, Hk, Hq // Hk, L, S)
-            s = mx.where(m, s, mx.finfo(s.dtype).min)
+        ke = k[:, :, None, :, :]
+        ve = v[:, :, None, :, :]
+        s = q @ ke.transpose(0, 1, 2, 4, 3)
         p = mx.softmax(s.astype(mx.float32), axis=-1).astype(s.dtype)
-        o = p @ v
-        return o.reshape(B, Hq, L, V)
-
-    for i in range(loops):
-        q = _sdpa(q, k, v)
-        q = upproject(q, w)
+        q = p @ ve
+        q = q.reshape(B, Hq, L, D)
     return q
 
 
-def sdpa(q, k, v, mask=None, w=None):
-    for i in range(loops):
-        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
-        q = upproject(q, w)
+def sdpa(q, k, v):
+    for _ in range(loops):
+        q = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=None)
     return q
 
 
-def time_self_attention_primitives():
-    mx.random.seed(3)
-    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
-    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
-    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mx.eval(q, k, v, w)
-    time_fn(attention, q, k, v, w=w)
+def quant_sdpa(q, k, v, bits=4):
+    for _ in range(loops):
+        q = mx.fast.quantized_scaled_dot_product_attention(
+            q, *k, *v, scale=1.0, mask=None, bits=bits
+        )
+    return q
 
 
-def time_self_attention_sdpa():
-    mx.random.seed(3)
-    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
-    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
-    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mx.eval(q, k, v, w)
-    time_fn(sdpa, q, k, v, w=w)
+def quant_attention(q, k, v, bits=4):
+    for _ in range(loops):
+        B, Hq, L, D = q.shape
+        Hk = k[0].shape[1]
+
+        q = q.reshape((B, Hk, Hq // Hk, L, D))
+        ke = tree_map(lambda x: mx.expand_dims(x, axis=2), k)
+        ve = tree_map(lambda x: mx.expand_dims(x, axis=2), v)
+
+        scores = mx.quantized_matmul(q, *ke, transpose=True, bits=bits)
+        scores = mx.softmax(scores, axis=-1)
+
+        q = mx.quantized_matmul(scores, *ve, transpose=False, bits=bits)
+        q = q.reshape((B, Hq, L, D))
+    return q
 
 
-def time_self_attention_sdpa_with_mask():
-    mx.random.seed(3)
-    q = mx.random.uniform(shape=(1, H, 1, D)).astype(dtype)
-    k = mx.random.uniform(shape=(1, H_k, L, D)).astype(dtype)
-    v = mx.random.uniform(shape=(1, H_k, L, V)).astype(dtype)
-    w = mx.random.uniform(shape=(D, V)).astype(dtype) if V != D else None
-    mask = mx.full((L,), True)
-    mask[L // 2 :] = False
-    mx.eval(q, k, v, mask, w)
+def time_self_attention_primitives(q, k, v):
+    time_fn(attention, q, k, v)
 
-    def sdpa_mask(*args):
-        return sdpa(*args, mask=mask, w=w)
 
-    def attention_mask(*args):
-        return attention(*args, mask=mask, w=w)
+def time_self_attention_sdpa(q, k, v):
+    time_fn(sdpa, q, k, v)
 
-    time_fn(attention_mask, q, k, v)
-    time_fn(sdpa_mask, q, k, v)
+
+def time_self_attention_quant_sdpa(q, k, v, bits=4):
+    time_fn(quant_sdpa, q, k, v, bits)
+
+
+def time_self_attention_quant_primitives(q, k, v, bits=4):
+    time_fn(quant_attention, q, k, v, bits)
 
 
 if __name__ == "__main__":
-    time_self_attention_sdpa()
-    time_self_attention_primitives()
-    time_self_attention_sdpa_with_mask()
+    mx.random.seed(3)
+    q = mx.random.uniform(shape=(1, H, 1, D), dtype=dtype)
+    k = mx.random.uniform(shape=(1, H_k, L, D), dtype=dtype)
+    v = mx.random.uniform(shape=(1, H_k, L, D), dtype=dtype)
+    mx.eval(q, k, v)
+
+    k_quant = mx.quantize(k, bits=bits)
+    v_quant = mx.quantize(v, bits=bits)
+    mx.eval(k_quant, v_quant)
+
+    k = mx.dequantize(*k_quant, bits=bits)
+    v = mx.dequantize(*v_quant, bits=bits)
+
+    time_self_attention_sdpa(q, k, v)
+    time_self_attention_quant_sdpa(q, k_quant, v_quant, bits)
+    time_self_attention_primitives(q, k, v)
+    time_self_attention_quant_primitives(q, k_quant, v_quant, bits)

benchmarks/python/single_ops.py

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

benchmarks/python/synchronize_bench.py

@@ -1,55 +0,0 @@
-import time
-
-import mlx.core as mx
-
-rank = mx.distributed.init().rank()
-
-
-def timeit(fn, a):
-
-    # warmup
-    for _ in range(5):
-        mx.eval(fn(a))
-
-    its = 10
-    tic = time.perf_counter()
-    for _ in range(its):
-        mx.eval(fn(a))
-    toc = time.perf_counter()
-    ms = 1000 * (toc - tic) / its
-    return ms
-
-
-def all_reduce_benchmark():
-    a = mx.ones((5, 5), mx.int32)
-
-    its_per_eval = 100
-
-    def fn(x):
-        for _ in range(its_per_eval):
-            x = mx.distributed.all_sum(x)
-            x = x - 1
-        return x
-
-    ms = timeit(fn, a) / its_per_eval
-    if rank == 0:
-        print(f"All Reduce: time per iteration {ms:.6f} (ms)")
-
-
-def all_gather_benchmark():
-    a = mx.ones((5, 5), mx.int32)
-    its_per_eval = 100
-
-    def fn(x):
-        for _ in range(its_per_eval):
-            x = mx.distributed.all_gather(x)[0]
-        return x
-
-    ms = timeit(fn, a) / its_per_eval
-    if rank == 0:
-        print(f"All gather: time per iteration {ms:.6f} (ms)")
-
-
-if __name__ == "__main__":
-    all_reduce_benchmark()
-    all_gather_benchmark()

cmake/FindNCCL.cmake

@@ -1,54 +0,0 @@
-# FindNCCL.cmake This module finds the NVIDIA NCCL library and its include
-# directories.
-
-set(NCCL_ROOT_DIR
-    $ENV{NCCL_ROOT_DIR}
-    CACHE PATH "Folder contains NVIDIA NCCL")
-
-find_path(
-  NCCL_INCLUDE_DIRS
-  NAMES nccl.h
-  HINTS ${NCCL_INCLUDE_DIR} ${NCCL_ROOT_DIR} ${NCCL_ROOT_DIR}/include
-        ${CUDA_TOOLKIT_ROOT_DIR}/include)
-
-if($ENV{USE_STATIC_NCCL})
-  message(
-    STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library")
-  set(NCCL_LIBNAME "libnccl_static.a")
-else()
-  set(NCCL_LIBNAME "nccl")
-endif()
-
-find_library(
-  NCCL_LIBRARIES
-  NAMES ${NCCL_LIBNAME}
-  HINTS ${NCCL_LIB_DIR}
-        ${NCCL_ROOT_DIR}
-        ${NCCL_ROOT_DIR}/lib
-        ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu
-        ${NCCL_ROOT_DIR}/lib64
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib
-        ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
-
-include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS
-                                  NCCL_LIBRARIES)
-
-if(NCCL_FOUND)
-  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h")
-  message(
-    STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}")
-  file(
-    STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED
-    REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$"
-    LIMIT_COUNT 1)
-  if(NCCL_MAJOR_VERSION_DEFINED)
-    string(REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" ""
-                         NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED})
-    message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}")
-  endif()
-  message(
-    STATUS
-      "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
-  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
-endif()

cmake/Findnvpl.cmake

@@ -1,3 +0,0 @@
-# This file does nothing but to suppress the cmake warning: "By not providing
-# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
-# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.

cmake/extension.cmake

@@ -1,7 +1,5 @@
 include(CMakeParseArguments)
 
-# clang format off
-#
 # ##############################################################################
 # Build metal library
 #
@@ -11,14 +9,11 @@ include(CMakeParseArguments)
 # Args: TARGET: Custom target to be added for the metal library TITLE: Name of
 # the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List
 # of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
-# files (like headers) DEBUG: Boolean, if true, enables debug compile options
-# for this specific library. If not provided, uses global MLX_METAL_DEBUG.
+# files (like headers)
 #
-# clang format on
-
 macro(mlx_build_metallib)
   # Parse args
-  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG)
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
   set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
   cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
@@ -26,11 +21,7 @@ macro(mlx_build_metallib)
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
   # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
-  if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
-    set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
-                               -frecord-sources)
-  endif()
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math)
 
   # Prepare metallib build command
   add_custom_command(

docs/Doxyfile

@@ -13,7 +13,7 @@ EXCLUDE_PATTERNS       = */private/*
 CREATE_SUBDIRS         = NO
 FULL_PATH_NAMES        = YES
 RECURSIVE              = YES
-GENERATE_HTML          = NO
+GENERATE_HTML          = YES
 GENERATE_LATEX         = NO
 GENERATE_XML           = YES
 XML_PROGRAMLISTING     = YES

docs/requirements.txt

@@ -1,5 +1,4 @@
 sphinx
 breathe
 sphinx-book-theme
-sphinx-copybutton
 mlx

docs/src/conf.py

@@ -10,7 +10,7 @@ import mlx.core as mx
 # -- Project information -----------------------------------------------------
 
 project = "MLX"
-copyright = "2023, Apple"
+copyright = "2023, MLX Contributors"
 author = "MLX Contributors"
 version = ".".join(mx.__version__.split(".")[:3])
 release = version
@@ -18,7 +18,6 @@ release = version
 # -- General configuration ---------------------------------------------------
 
 extensions = [
-    "sphinx_copybutton",
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",
     "sphinx.ext.intersphinx",

docs/src/dev/custom_metal_kernels.rst

@@ -8,12 +8,11 @@ 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
 
+  def exp_elementwise(a: mx.array):
       source = """
           uint elem = thread_position_in_grid.x;
           T tmp = inp[elem];
@@ -26,8 +25,6 @@ Let's write a custom kernel that computes ``exp`` elementwise:
           output_names=["out"],
           source=source,
       )
-
-  def exp_elementwise(a: mx.array):
       outputs = kernel(
           inputs=[a],
           template=[("T", mx.float32)],
@@ -42,13 +39,8 @@ Let's write a custom kernel that computes ``exp`` elementwise:
   b = exp_elementwise(a)
   assert mx.allclose(b, mx.exp(a))
 
-Every time you make a kernel, a new Metal library is created and possibly
-JIT compiled. To reduce the overhead from that, build the kernel once with
-:func:`fast.metal_kernel` and then use it many times.
-
 .. note::
-   Only pass the body of the Metal kernel in ``source``. The function
-   signature is generated automatically.
+    We are only required to pass the body of the Metal kernel in ``source``.
 
 The full function signature will be generated using:
 
@@ -86,34 +78,29 @@ Putting this all together, the generated function signature for ``myexp`` is as
 
   template [[host_name("custom_kernel_myexp_float")]] [[kernel]] decltype(custom_kernel_myexp_float<float>) custom_kernel_myexp_float<float>;
 
-Note: ``grid`` and ``threadgroup`` are parameters to the Metal `dispatchThreads
-<https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads>`_
-function. This means we will launch ``mx.prod(grid)`` threads, subdivided into
-``threadgroup`` size threadgroups.  For optimal performance, each thread group
-dimension should be less than or equal to the corresponding grid dimension.
+Note: ``grid`` and ``threadgroup`` are parameters to the Metal `dispatchThreads <https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads>`_ function.
+This means we will launch ``mx.prod(grid)`` threads, subdivided into ``threadgroup`` size threadgroups.
+For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension.
 
-Passing ``verbose=True`` to :func:`ast.metal_kernel.__call__` will print the
-generated code for debugging purposes.
+Passing ``verbose=True`` to ``mx.fast.metal_kernel.__call__`` will print the generated code for debugging purposes.
 
 Using Shape/Strides
 -------------------
 
-:func:`fast.metal_kernel` supports an argument ``ensure_row_contiguous`` which
-is ``True`` by default. This will copy the array inputs if needed
-before the kernel is launched to ensure that the memory layout is row
-contiguous.  Generally this makes writing the kernel easier, since we don't
-have to worry about gaps or the ordering of the dims when indexing.
+``mx.fast.metal_kernel`` supports an argument ``ensure_row_contiguous`` which is ``True`` by default.
+This will copy the ``mx.array`` inputs if needed before the kernel is launched to ensure that the memory layout is row contiguous.
+Generally this makes writing the kernel easier, since we don't have to worry about gaps or the ordering of the dims
+when indexing.
 
-If we want to avoid this copy, :func:`fast.metal_kernel` automatically passes
-``a_shape``, ``a_strides`` and ``a_ndim`` for each input array ``a`` if any are
-present in ``source``. We can then use MLX's built in indexing utils to fetch
-the right elements for each thread.
+If we want to avoid this copy, ``metal_kernel`` automatically passes ``a_shape``, ``a_strides`` and ``a_ndim`` for each
+input array ``a`` if any are present in ``source``.
+We can then use MLX's built in indexing utils to fetch the right elements for each thread.
 
-Let's convert ``myexp`` above to support arbitrarily strided arrays without
-relying on a copy from ``ensure_row_contiguous``:
+Let's convert ``myexp`` above to support arbitrarily strided arrays without relying on a copy from ``ensure_row_contiguous``:
 
 .. code-block:: python
 
+  def exp_elementwise(a: mx.array):
       source = """
           uint elem = thread_position_in_grid.x;
           // Utils from `mlx/backend/metal/kernels/utils.h` are automatically included
@@ -127,11 +114,8 @@ relying on a copy from ``ensure_row_contiguous``:
           name="myexp_strided",
           input_names=["inp"],
           output_names=["out"],
-      source=source,
-      ensure_row_contiguous=False,
+          source=source
       )
-
-  def exp_elementwise(a: mx.array):
       outputs = kernel(
           inputs=[a],
           template=[("T", mx.float32)],
@@ -139,6 +123,7 @@ relying on a copy from ``ensure_row_contiguous``:
           threadgroup=(256, 1, 1),
           output_shapes=[a.shape],
           output_dtypes=[a.dtype],
+          ensure_row_contiguous=False,
       )
       return outputs[0]
 
@@ -198,13 +183,25 @@ We'll start with the following MLX implementation using standard ops:
 
         return output
 
-Now let's use :func:`custom_function` together with :func:`fast.metal_kernel`
+Now let's use ``mx.custom_function`` together with ``mx.fast.metal_kernel``
 to write a fast GPU kernel for both the forward and backward passes.
 
 First we'll implement the forward pass as a fused kernel:
 
 .. code-block:: python
 
+    @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."
+
         source = """
             uint elem = thread_position_in_grid.x;
             int H = x_shape[1];
@@ -254,26 +251,12 @@ First we'll implement the forward pass as a fused kernel:
 
             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)],
@@ -298,11 +281,11 @@ On an M1 Max, we see a big performance improvement:
 Grid Sample VJP
 ---------------
 
-Since we decorated ``grid_sample`` with :func:`custom_function`, we can now
-define its custom vjp transform so MLX can differentiate it.
+Since we decorated ``grid_sample`` with ``mx.custom_function``, we can now define
+its custom vjp transform so MLX can differentiate it.
 
 The backwards pass requires atomically updating ``x_grad``/``grid_grad`` and so
-requires a few extra :func:`fast.metal_kernel` features:
+requires a few extra ``mx.fast.metal_kernel`` features:
 
 * ``init_value=0``
     Initialize all of the kernel's outputs to this value before it runs. This allows us to update only part of the output arrays with the kernel.
@@ -316,6 +299,14 @@ We can then implement the backwards pass as follows:
 
 .. code-block:: python
 
+    @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."
+
         source = """
             uint elem = thread_position_in_grid.x;
             int H = x_shape[1];
@@ -415,15 +406,6 @@ We can then implement the backwards pass as follows:
             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

docs/src/dev/extensions.rst

@@ -22,12 +22,12 @@ You can do that in MLX directly:
 This function performs that operation while leaving the implementation and
 function transformations to MLX.
 
-However, you may want to customize the underlying implementation, perhaps to
-make it faster. In this tutorial we will go through adding custom extensions.
-It will cover:
+However you may need to customize the underlying implementation, perhaps to
+make it faster or for custom differentiation. In this tutorial we will go
+through adding custom extensions. It will cover:
 
 * The structure of the MLX library.
-* Implementing a CPU operation.
+* Implementing a CPU operation that redirects to Accelerate_ when appropriate.
 * Implementing a GPU operation using metal.
 * Adding the ``vjp`` and ``jvp`` function transformation.
 * Building a custom extension and binding it to python.
@@ -45,7 +45,7 @@ Operations
 Operations are the front-end functions that operate on arrays. They are defined
 in the C++ API (:ref:`cpp_ops`), and the Python API (:ref:`ops`) binds them.
 
-We would like an operation :meth:`axpby` that takes in two arrays, ``x`` and
+We would like an operation, :meth:`axpby` that takes in two arrays ``x`` and
 ``y``, and two scalars, ``alpha`` and ``beta``. This is how to define it in
 C++:
 
@@ -55,7 +55,7 @@ C++:
     *  Scale and sum two vectors element-wise
     *  z = alpha * x + beta * y
     *
-    *  Use NumPy-style broadcasting between x and y
+    *  Follow numpy style broadcasting between x and y
     *  Inputs are upcasted to floats if needed
     **/
     array axpby(
@@ -66,7 +66,7 @@ C++:
         StreamOrDevice s = {} // Stream on which to schedule the operation
     );
 
-The simplest way to implement this is with existing operations:
+The simplest way to this operation is in terms of existing operations:
 
 .. code-block:: C++
 
@@ -93,9 +93,9 @@ Primitives
 ^^^^^^^^^^^
 
 A :class:`Primitive` is part of the computation graph of an :class:`array`. It
-defines how to create output arrays given input arrays. Further, a
+defines how to create outputs arrays given a input arrays. Further, a
 :class:`Primitive` has methods to run on the CPU or GPU and for function
-transformations such as ``vjp`` and ``jvp``.  Let's go back to our example to be
+transformations such as ``vjp`` and ``jvp``.  Lets go back to our example to be
 more concrete:
 
 .. code-block:: C++
@@ -128,7 +128,7 @@ more concrete:
         /** The vector-Jacobian product. */
         std::vector<array> vjp(
             const std::vector<array>& primals,
-            const std::vector<array>& cotangents,
+            const array& cotan,
             const std::vector<int>& argnums,
             const std::vector<array>& outputs) override;
 
@@ -138,13 +138,13 @@ more concrete:
         * representing the vectorized computation and the axis which
         * corresponds to the output vectorized dimension.
         */
-        std::pair<std::vector<array>, std::vector<int>> vmap(
+        virtual std::pair<std::vector<array>, std::vector<int>> vmap(
             const std::vector<array>& inputs,
             const std::vector<int>& axes) override;
 
-        /** The name of primitive. */
-        const char* name() const override {
-          return "Axpby";
+        /** Print the primitive. */
+        void print(std::ostream& os) override {
+            os << "Axpby";
         }
 
         /** Equivalence check **/
@@ -153,6 +153,9 @@ more concrete:
       private:
         float alpha_;
         float beta_;
+
+        /** Fall back implementation for evaluation on CPU */
+        void eval(const std::vector<array>& inputs, array& out);
     };
 
 The :class:`Axpby` class derives from the base :class:`Primitive` class. The
@@ -185,7 +188,7 @@ Let's reimplement our operation now in terms of our :class:`Axpby` primitive.
         auto promoted_dtype = promote_types(x.dtype(), y.dtype());
 
         // Upcast to float32 for non-floating point inputs x and y
-        auto out_dtype = issubdtype(promoted_dtype, float32)
+        auto out_dtype = is_floating_point(promoted_dtype)
             ? promoted_dtype
             : promote_types(promoted_dtype, float32);
 
@@ -231,9 +234,11 @@ the execution of the computation graph, and calls :meth:`Axpby::eval_cpu` or
 Implementing the CPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Let's start by implementing :meth:`Axpby::eval_cpu`.
+Let's start by implementing a naive and generic version of
+:meth:`Axpby::eval_cpu`. We declared this as a private member function of
+:class:`Axpby` earlier called :meth:`Axpby::eval`.
 
-The method will go over each element of the output array, find the
+Our naive method will go over each element of the output array, find the
 corresponding input elements of ``x`` and ``y`` and perform the operation
 point-wise. This is captured in the templated function :meth:`axpby_impl`.
 
@@ -241,46 +246,36 @@ point-wise. This is captured in the templated function :meth:`axpby_impl`.
 
     template <typename T>
     void axpby_impl(
-      const mx::array& x,
-      const mx::array& y,
-      mx::array& out,
+            const array& x,
+            const array& y,
+            array& out,
             float alpha_,
-      float beta_,
-      mx::Stream stream) {
-    out.set_data(mx::allocator::malloc(out.nbytes()));
+            float beta_) {
+        // We only allocate memory when we are ready to fill the output
+        // malloc_or_wait synchronously allocates available memory
+        // There may be a wait executed here if the allocation is requested
+        // under memory-pressured conditions
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
-    // Get the CPU command encoder and register input and output arrays
-    auto& encoder = mx::cpu::get_command_encoder(stream);
-    encoder.set_input_array(x);
-    encoder.set_input_array(y);
-    encoder.set_output_array(out);
-
-    // 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_]() {
+        // Collect input and output data pointers
+        const T* x_ptr = x.data<T>();
+        const T* y_ptr = y.data<T>();
+        T* out_ptr = out.data<T>();
 
         // 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++) {
+        for (size_t out_idx = 0; out_idx < out.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);
+            auto x_offset = elem_to_loc(out_idx, x.shape(), x.strides());
+            auto y_offset = elem_to_loc(out_idx, y.shape(), y.strides());
 
             // We allocate the output to be contiguous and regularly strided
             // (defaults to row major) and hence it doesn't need additional mapping
             out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
         }
-    });
     }
 
 Our implementation should work for all incoming floating point arrays.
@@ -289,32 +284,112 @@ Accordingly, we add dispatches for ``float32``, ``float16``, ``bfloat16`` and
 
 .. code-block:: C++
 
-    void Axpby::eval_cpu(
-        const std::vector<mx::array>& inputs,
-        std::vector<mx::array>& outputs) {
+    /** Fall back implementation for evaluation on CPU */
+    void Axpby::eval(
+      const std::vector<array>& inputs,
+      const std::vector<array>& outputs) {
         auto& x = inputs[0];
         auto& y = inputs[1];
         auto& out = outputs[0];
 
         // Dispatch to the correct dtype
-      if (out.dtype() == mx::float32) {
-        return axpby_impl<float>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::float16) {
-        return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::bfloat16) {
-        return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_, stream());
-      } else if (out.dtype() == mx::complex64) {
-        return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_, stream());
+        if (out.dtype() == float32) {
+            return axpby_impl<float>(x, y, out, alpha_, beta_);
+        } else if (out.dtype() == float16) {
+            return axpby_impl<float16_t>(x, y, out, alpha_, beta_);
+        } else if (out.dtype() == bfloat16) {
+            return axpby_impl<bfloat16_t>(x, y, out, alpha_, beta_);
+        } else if (out.dtype() == complex64) {
+            return axpby_impl<complex64_t>(x, y, out, alpha_, beta_);
         } else {
             throw std::runtime_error(
-            "Axpby is only supported for floating point types.");
+                "[Axpby] Only supports floating point types.");
         }
     }
 
+This is good as a fallback implementation. We can use the ``axpby`` routine
+provided by the Accelerate_ framework for a faster implementation in certain
+cases:
+
+#.  Accelerate does not provide implementations of ``axpby`` for half precision
+    floats. We can only use it for ``float32`` types.
+#.  Accelerate assumes the inputs ``x`` and ``y`` are contiguous and all
+    elements have fixed strides between them. We only direct to Accelerate
+    if both ``x`` and ``y`` are row contiguous or column contiguous.
+#.  Accelerate performs the routine ``Y = (alpha * X) + (beta * Y)`` in-place.
+    MLX expects to write the output to a new array. We must copy the elements
+    of ``y`` into the output and use that as an input to ``axpby``.
+
+Let's write an implementation that uses Accelerate in the right conditions.
+It allocates data for the output, copies ``y`` into it, and then calls the
+:func:`catlas_saxpby` from accelerate.
+
+.. code-block:: C++
+
+    template <typename T>
+    void axpby_impl_accelerate(
+            const array& x,
+            const array& y,
+            array& out,
+            float alpha_,
+            float beta_) {
+        // Accelerate library provides catlas_saxpby which does
+        // Y = (alpha * X) + (beta * Y) in place
+        // To use it, we first copy the data in y over to the output array
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
+
+        // We then copy over the elements using the contiguous vector specialization
+        copy_inplace(y, out, CopyType::Vector);
+
+        // Get x and y pointers for catlas_saxpby
+        const T* x_ptr = x.data<T>();
+        T* y_ptr = out.data<T>();
+
+        T alpha = static_cast<T>(alpha_);
+        T beta = static_cast<T>(beta_);
+
+        // Call the inplace accelerate operator
+        catlas_saxpby(
+            /* N = */ out.size(),
+            /* ALPHA = */ alpha,
+            /* X = */ x_ptr,
+            /* INCX = */ 1,
+            /* BETA = */ beta,
+            /* Y = */ y_ptr,
+            /* INCY = */ 1);
+    }
+
+For inputs that do not fit the criteria for accelerate, we fall back to
+:meth:`Axpby::eval`. With this in mind, let's finish our
+:meth:`Axpby::eval_cpu`.
+
+.. code-block:: C++
+
+    /** Evaluate primitive on CPU using accelerate specializations */
+    void Axpby::eval_cpu(
+      const std::vector<array>& inputs,
+      const std::vector<array>& outputs) {
+        assert(inputs.size() == 2);
+        auto& x = inputs[0];
+        auto& y = inputs[1];
+        auto& out = outputs[0];
+
+        // Accelerate specialization for contiguous single precision float arrays
+        if (out.dtype() == float32 &&
+            ((x.flags().row_contiguous && y.flags().row_contiguous) ||
+            (x.flags().col_contiguous && y.flags().col_contiguous))) {
+            axpby_impl_accelerate<float>(x, y, out, alpha_, beta_);
+            return;
+        }
+
+        // Fall back to common back-end if specializations are not available
+        eval(inputs, outputs);
+    }
+
 Just this much is enough to run the operation :meth:`axpby` on a CPU stream! If
 you do not plan on running the operation on the GPU or using transforms on
 computation graphs that contain :class:`Axpby`, you can stop implementing the
-primitive here.
+primitive here and enjoy the speed-ups you get from the Accelerate library.
 
 Implementing the GPU Back-end
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -345,8 +420,8 @@ element in the output.
             constant const float& alpha [[buffer(3)]],
             constant const float& beta [[buffer(4)]],
             constant const int* shape [[buffer(5)]],
-            constant const int64_t* x_strides [[buffer(6)]],
-            constant const int64_t* y_strides [[buffer(7)]],
+            constant const size_t* x_strides [[buffer(6)]],
+            constant const size_t* y_strides [[buffer(7)]],
             constant const int& ndim [[buffer(8)]],
             uint index [[thread_position_in_grid]]) {
         // Convert linear indices to offsets in array
@@ -363,10 +438,24 @@ each instantiation a unique host name so we can identify it.
 
 .. code-block:: C++
 
-    instantiate_kernel("axpby_general_float32", axpby_general, float)
-    instantiate_kernel("axpby_general_float16", axpby_general, float16_t)
-    instantiate_kernel("axpby_general_bfloat16", axpby_general, bfloat16_t)
-    instantiate_kernel("axpby_general_complex64", axpby_general, complex64_t)
+    #define instantiate_axpby(type_name, type)              \
+        template [[host_name("axpby_general_" #type_name)]] \
+        [[kernel]] void axpby_general<type>(                \
+            device const type* x [[buffer(0)]],             \
+            device const type* y [[buffer(1)]],             \
+            device type* out [[buffer(2)]],                 \
+            constant const float& alpha [[buffer(3)]],      \
+            constant const float& beta [[buffer(4)]],       \
+            constant const int* shape [[buffer(5)]],        \
+            constant const size_t* x_strides [[buffer(6)]], \
+            constant const size_t* y_strides [[buffer(7)]], \
+            constant const int& ndim [[buffer(8)]],         \
+            uint index [[thread_position_in_grid]]);
+
+    instantiate_axpby(float32, float);
+    instantiate_axpby(float16, half);
+    instantiate_axpby(bfloat16, bfloat16_t);
+    instantiate_axpby(complex64, complex64_t);
 
 The logic to determine the kernel, set the inputs, resolve the grid dimensions,
 and dispatch to the GPU are contained in :meth:`Axpby::eval_gpu` as shown
@@ -391,17 +480,17 @@ below.
         auto& d = metal::device(s.device);
 
         // Allocate output memory
-        out.set_data(allocator::malloc(out.nbytes()));
+        out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
         // Resolve name of kernel
-        std::stream kname;
-        kname = "axpby_general_" + type_to_name(out);
+        std::ostringstream kname;
+        kname << "axpby_" << "general_" << type_to_name(out);
 
-        // Load the metal library
-        auto lib = d.get_library("mlx_ext", current_binary_dir());
+        // Make sure the metal library is available
+        d.register_library("mlx_ext");
 
         // Make a kernel from this metal library
-        auto kernel = d.get_kernel(kname, lib);
+        auto kernel = d.get_kernel(kname.str(), "mlx_ext");
 
         // Prepare to encode kernel
         auto& compute_encoder = d.get_command_encoder(s.index);
@@ -469,7 +558,7 @@ one we just defined:
             const std::vector<array>& tangents,
             const std::vector<int>& argnums) {
         // Forward mode diff that pushes along the tangents
-        // The jvp transform on the primitive can be built with ops
+        // The jvp transform on the primitive can built with ops
         // that are scheduled on the same stream as the primitive
 
         // If argnums = {0}, we only push along x in which case the
@@ -481,7 +570,7 @@ one we just defined:
             auto scale_arr = array(scale, tangents[0].dtype());
             return {multiply(scale_arr, tangents[0], stream())};
         }
-        // If argnums = {0, 1}, we take contributions from both
+        // If, argnums = {0, 1}, we take contributions from both
         // which gives us jvp = tangent_x * alpha + tangent_y * beta
         else {
             return {axpby(tangents[0], tangents[1], alpha_, beta_, stream())};
@@ -735,7 +824,7 @@ Let's look at a simple script and its results:
 
     print(f"c shape: {c.shape}")
     print(f"c dtype: {c.dtype}")
-    print(f"c is correct: {mx.all(c == 6.0).item()}")
+    print(f"c correct: {mx.all(c == 6.0).item()}")
 
 Output:
 
@@ -743,13 +832,13 @@ Output:
 
     c shape: [3, 4]
     c dtype: float32
-    c is correct: True
+    c correctness: True
 
 Results
 ^^^^^^^
 
 Let's run a quick benchmark and see how our new ``axpby`` operation compares
-with the naive :meth:`simple_axpby` we first defined.
+with the naive :meth:`simple_axpby` we first defined on the CPU.
 
 .. code-block:: python
 
@@ -757,11 +846,13 @@ with the naive :meth:`simple_axpby` we first defined.
     from mlx_sample_extensions import axpby
     import time
 
+    mx.set_default_device(mx.cpu)
+
     def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
         return alpha * x + beta * y
 
-    M = 4096
-    N = 4096
+    M = 256
+    N = 512
 
     x = mx.random.normal((M, N))
     y = mx.random.normal((M, N))
@@ -772,24 +863,24 @@ with the naive :meth:`simple_axpby` we first defined.
 
     def bench(f):
         # Warm up
-        for i in range(5):
+        for i in range(100):
             z = f(x, y, alpha, beta)
             mx.eval(z)
 
         # Timed run
         s = time.time()
-        for i in range(100):
+        for i in range(5000):
             z = f(x, y, alpha, beta)
             mx.eval(z)
         e = time.time()
-        return 1000 * (e - s) / 100
+        return e - s
 
     simple_time = bench(simple_axpby)
     custom_time = bench(axpby)
 
-    print(f"Simple axpby: {simple_time:.3f} ms | Custom axpby: {custom_time:.3f} ms")
+    print(f"Simple axpby: {simple_time:.3f} s | Custom axpby: {custom_time:.3f} s")
 
-The results are ``Simple axpby: 1.559 ms | Custom axpby: 0.774 ms``. We see
+The results are ``Simple axpby: 0.114 s | Custom axpby: 0.109 s``. We see
 modest improvements right away!
 
 This operation is now good to be used to build other operations, in

docs/src/dev/mlx_in_cpp.rst

@@ -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

docs/src/index.rst

@@ -45,7 +45,6 @@ are the CPU and GPU.
    usage/numpy
    usage/distributed
    usage/using_streams
-   usage/export
 
 .. toctree::
    :caption: Examples
@@ -62,7 +61,6 @@ are the CPU and GPU.
    python/array
    python/data_types
    python/devices_and_streams
-   python/export
    python/ops
    python/random
    python/transforms
@@ -70,8 +68,6 @@ are the CPU and GPU.
    python/fft
    python/linalg
    python/metal
-   python/cuda
-   python/memory_management
    python/nn
    python/optimizers
    python/distributed
@@ -90,4 +86,3 @@ are the CPU and GPU.
    dev/extensions
    dev/metal_debugger
    dev/custom_metal_kernels
-   dev/mlx_in_cpp

docs/src/install.rst

@@ -1,5 +1,3 @@
-.. _build_and_install:
-
 Build and Install
 =================
 
@@ -13,48 +11,22 @@ silicon computer is
 
     pip install mlx
 
-To install from PyPI your system must meet the following requirements:
+To install from PyPI you must meet the following requirements:
 
 - Using an M series chip (Apple silicon)
-- Using a native Python >= 3.10
-- macOS >= 14.0
+- Using a native Python >= 3.9
+- macOS >= 13.5
 
 .. note::
-    MLX is only available on devices running macOS >= 14.0 and higher.
+    MLX is only available on devices running macOS >= 13.5
+    It is highly recommended to use macOS 14 (Sonoma)
 
-CUDA
-^^^^
 
-MLX has a CUDA backend which you can install with:
+MLX is also available on conda-forge. To install MLX with conda do:
 
 .. code-block:: shell
 
-    pip install mlx[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.10
-
-
-CPU-only (Linux)
-^^^^^^^^^^^^^^^^
-
-For a CPU-only version of MLX that runs on Linux use:
-
-.. code-block:: shell
-
-    pip install mlx[cpu]
-
-To install the CPU-only package from PyPi your system must meet the following
-requirements:
-
-- Linux distribution with glibc >= 2.35
-- Python >= 3.10
+   conda install conda-forge::mlx
 
 
 Troubleshooting
@@ -81,7 +53,7 @@ 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``
+- `cmake <https://cmake.org/>`_ -- version 3.24 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0
 
 .. note::
@@ -91,8 +63,6 @@ Build Requirements
 Python API
 ^^^^^^^^^^
 
-.. _python install:
-
 To build and install the MLX python library from source, first, clone MLX from
 `its GitHub repo <https://github.com/ml-explore/mlx>`_:
 
@@ -104,20 +74,20 @@ Then simply build and install MLX using pip:
 
 .. code-block:: shell
 
-  pip install .
+  CMAKE_BUILD_PARALLEL_LEVEL=8 pip install .
 
 For developing, install the package with development dependencies, and use an
 editable install:
 
 .. code-block:: shell
 
-  pip install -e ".[dev]"
+  CMAKE_BUILD_PARALLEL_LEVEL=8 pip install -e ".[dev]"
 
 Once the development dependencies are installed, you can build faster with:
 
 .. code-block:: shell
 
- python setup.py build_ext --inplace
+ CMAKE_BUILD_PARALLEL_LEVEL=8 python setup.py build_ext --inplace
 
 Run the tests with:
 
@@ -135,8 +105,6 @@ IDE:
 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
@@ -215,7 +183,6 @@ should point to the path to the built metal library.
 
       xcrun -sdk macosx --show-sdk-version
 
-
 Binary Size Minimization
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -244,50 +211,6 @@ be anwywhere from a few hundred millisecond to a few seconds depending on the
 application. Once a kernel is compiled, it will be cached by the system. The
 Metal kernel cache persists across reboots.
 
-Linux
-^^^^^
-
-To build from source on Linux (CPU only), install the BLAS and LAPACK headers.
-For example on Ubuntu, run the following:
-
-.. code-block:: shell
-
-   apt-get update -y
-   apt-get install libblas-dev liblapack-dev liblapacke-dev -y
-
-From here follow the instructions to install either the :ref:`Python <python
-install>` or :ref:`C++ <cpp install>` APIs.
-
-CUDA
-^^^^
-
-To build from source on Linux with CUDA, install the BLAS and LAPACK headers
-and the CUDA toolkit. For example on Ubuntu, run the following:
-
-.. code-block:: shell
-
-   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
-   dpkg -i cuda-keyring_1.1-1_all.deb
-   apt-get update -y
-   apt-get -y install cuda-toolkit-12-9
-   apt-get install libblas-dev liblapack-dev liblapacke-dev libcudnn9-dev-cuda-12 -y
-
-
-When building either the Python or C++ APIs make sure to pass the cmake flag
-``MLX_BUILD_CUDA=ON``. For example, to build the Python API run:
-
-.. code-block:: shell
-
-  CMAKE_ARGS="-DMLX_BUILD_CUDA=ON" pip install -e ".[dev]"
-
-To build the C++ package run:
-
-.. code-block:: shell
-
-   mkdir -p build && cd build
-   cmake .. -DMLX_BUILD_CUDA=ON && make -j
-
-
 Troubleshooting
 ^^^^^^^^^^^^^^^
 

docs/src/python/array.rst

@@ -19,8 +19,6 @@ Array
     array.ndim
     array.shape
     array.size
-    array.real
-    array.imag
     array.abs
     array.all
     array.any
@@ -40,7 +38,6 @@ Array
     array.log10
     array.log1p
     array.log2
-    array.logcumsumexp
     array.logsumexp
     array.max
     array.mean

docs/src/python/cuda.rst

@@ -1,9 +0,0 @@
-CUDA
-=====
-
-.. currentmodule:: mlx.core.cuda
-
-.. autosummary::
-  :toctree: _autosummary
-
-  is_available

docs/src/python/data_types.rst

@@ -51,20 +51,11 @@ The default floating point type is ``float32`` and the default integer type is
    * - ``float32``
      - 4 
      - 32-bit float
-   * - ``float64``
-     - 4 
-     - 64-bit double
    * - ``complex64``
      - 8 
      - 64-bit complex float
 
 
-.. note::
-
-    Arrays with type ``float64`` only work with CPU operations. Using
-    ``float64`` arrays on the GPU will result in an exception.
-
-
 Data type are aranged in a hierarchy. See the :obj:`DtypeCategory` object
 documentation for more information. Use :func:`issubdtype` to determine if one
 ``dtype`` (or category) is a subtype of another category.
@@ -75,4 +66,3 @@ documentation for more information. Use :func:`issubdtype` to determine if one
    Dtype
    DtypeCategory
    issubdtype
-   finfo

docs/src/python/export.rst

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

docs/src/python/fast.rst

@@ -13,4 +13,3 @@ Fast
   rope
   scaled_dot_product_attention
   metal_kernel
-  cuda_kernel

docs/src/python/fft.rst

@@ -20,5 +20,3 @@ FFT
   irfft2
   rfftn
   irfftn
-  fftshift
-  ifftshift

docs/src/python/linalg.rst

@@ -16,12 +16,5 @@ Linear Algebra
     cross
     qr
     svd
-    eigvals
-    eig
     eigvalsh
     eigh
-    lu
-    lu_factor
-    pinv
-    solve
-    solve_triangular

docs/src/python/memory_management.rst

@@ -1,16 +0,0 @@
-Memory Management
-=================
-
-.. currentmodule:: mlx.core
-
-.. autosummary::
-  :toctree: _autosummary
-
-  get_active_memory
-  get_peak_memory
-  reset_peak_memory
-  get_cache_memory
-  set_memory_limit
-  set_cache_limit
-  set_wired_limit
-  clear_cache

docs/src/python/metal.rst

@@ -8,5 +8,13 @@ Metal
 
   is_available
   device_info
+  get_active_memory
+  get_peak_memory
+  reset_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit
+  set_wired_limit
+  clear_cache
   start_capture
   stop_capture

docs/src/python/nn.rst

@@ -174,7 +174,6 @@ In detail:
 
    value_and_grad
    quantize
-   average_gradients
 
 .. toctree::
 

docs/src/python/nn/functions.rst

@@ -27,7 +27,6 @@ simple functions.
    mish
    prelu
    relu
-   relu2
    relu6
    selu
    sigmoid

docs/src/python/nn/layers.rst

@@ -50,7 +50,6 @@ Layers
    QuantizedLinear
    RMSNorm
    ReLU
-   ReLU2
    ReLU6
    RNN
    RoPE

docs/src/python/ops.rst

@@ -32,16 +32,13 @@ Operations
    atleast_2d
    atleast_3d
    bitwise_and
-   bitwise_invert
    bitwise_or
    bitwise_xor
    block_masked_mm
-   broadcast_arrays
    broadcast_to
    ceil
    clip
    concatenate
-   contiguous
    conj
    conjugate
    convolve
@@ -92,7 +89,6 @@ Operations
    isneginf
    isposinf
    issubdtype
-   kron
    left_shift
    less
    less_equal
@@ -103,7 +99,6 @@ Operations
    log10
    log1p
    logaddexp
-   logcumsumexp
    logical_not
    logical_and
    logical_or
@@ -112,7 +107,6 @@ Operations
    max
    maximum
    mean
-   median
    meshgrid
    min
    minimum
@@ -150,8 +144,6 @@ Operations
    sign
    sin
    sinh
-   slice
-   slice_update
    softmax
    sort
    split
@@ -176,7 +168,6 @@ Operations
    tri
    tril
    triu
-   unflatten
    var
    view
    where

docs/src/python/optimizers.rst

@@ -51,14 +51,14 @@ the saved state. Here's a simple example:
    optimizer.update(model, grads)
 
    # Save the state
-   state = tree_flatten(optimizer.state, destination={})
-   mx.save_safetensors("optimizer.safetensors", state)
+   state = tree_flatten(optimizer.state)
+   mx.save_safetensors("optimizer.safetensors", dict(state))
 
    # Later on, for example when loading from a checkpoint,
    # recreate the optimizer and load the state
    optimizer = optim.Adam(learning_rate=1e-2)
 
-   state = tree_unflatten(mx.load("optimizer.safetensors"))
+   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
    optimizer.state = state
 
 Note, not every optimizer configuation parameter is saved in the state. For

docs/src/python/optimizers/common_optimizers.rst

@@ -18,5 +18,3 @@ Common Optimizers
    AdamW
    Adamax
    Lion
-   MultiOptimizer
-   Muon

docs/src/python/transforms.rst

@@ -9,7 +9,6 @@ Transforms
   :toctree: _autosummary
 
    eval
-   async_eval
    compile
    custom_function
    disable_compile

docs/src/usage/compile.rst

@@ -130,8 +130,8 @@ Now make an array, and benchmark both functions:
 .. code-block:: python
 
   x = mx.random.uniform(shape=(32, 1000, 4096))
-  timeit(gelu, x)
-  timeit(mx.compile(gelu), x)
+  timeit(nn.gelu, x)
+  timeit(mx.compile(nn.gelu), x)
 
 On an M1 Max the times are 15.5 and 3.1 milliseconds. The compiled ``gelu`` is
 five times faster.
@@ -225,7 +225,7 @@ In some cases returning updated state can be pretty inconvenient. Hence,
   def fun(x, y):
       z = x + y
       state.append(z)
-      return mx.exp(z)
+      return mx.exp(z), state
 
   fun(mx.array(1.0), mx.array(2.0))
   # Prints [array(3, dtype=float32)]
@@ -421,77 +421,3 @@ the most opportunity to optimize the computation graph:
   # Compiling the outer function is good to do as it will likely
   # be faster even though the inner functions are compiled
   fun = mx.compile(outer)
-
-
-
-.. _shapeless_compile:
-
-Shapeless Compilation
----------------------
-
-When the shape of an input to a compiled function changes, the function is
-recompiled. You can compile a function once and run it on inputs with
-variable shapes by specifying ``shapeless=True`` to :func:`compile`. In this
-case changes to the shapes of the inputs do not cause the function to be
-recompiled.
-
-.. code-block:: python
-
-  def fun(x, y):
-      return mx.abs(x + y)
-
-  compiled_fun = mx.compile(fun, shapeless=True)
-
-  x = mx.array(1.0)
-  y = mx.array(-2.0)
-
-  # Firt call compiles the function
-  print(compiled_fun(x, y))
-
-  # Second call with different shapes
-  # does not recompile the function
-  x = mx.array([1.0, -6.0])
-  y = mx.array([-2.0, 3.0])
-  print(compiled_fun(x, y))
-
-
-Use shapeless compilations carefully. Since compilation is not triggered when
-shapes change, any graphs which are conditional on the input shapes will not
-work as expected. Shape-dependent computations are common and sometimes subtle
-to detect. For example:
-
-.. code-block:: python
-
-  def fun(x):
-      return x.reshape(x.shape[0] * x.shape[1], -1)
-
-  compiled_fun = mx.compile(fun, shapeless=True)
-
-  x = mx.random.uniform(shape=(2, 3, 4))
-
-  out = compiled_fun(x)
-
-  x = mx.random.uniform(shape=(5, 5, 3))
-
-  # Error, can't reshape (5, 5, 3) to (6, -1)
-  out = compiled_fun(x)
-
-The second call to the ``compiled_fun`` fails because of the call to
-:func:`reshape` which uses the static shape of ``x`` in the first call. We can
-fix this by using :func:`flatten` to avoid hardcoding the shape of ``x``:
-
-.. code-block:: python
-
-  def fun(x):
-      return x.flatten(0, 1)
-
-  compiled_fun = mx.compile(fun, shapeless=True)
-
-  x = mx.random.uniform(shape=(2, 3, 4))
-
-  out = compiled_fun(x)
-
-  x = mx.random.uniform(shape=(5, 5, 3))
-
-  # Ok
-  out = compiled_fun(x)

docs/src/usage/distributed.rst

@@ -5,28 +5,21 @@ Distributed Communication
 
 .. currentmodule:: mlx.core.distributed
 
-MLX supports distributed communication operations that allow the computational cost
-of training or inference to be shared across many physical machines. At the
-moment we support three different communication backends:
-
-* `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ a
-  full-featured and mature distributed communications library
-* A **ring** backend of our own that uses native TCP sockets. It should be
-  faster for thunderbolt connections, but it also works over Ethernet.
-* `nccl <https://developer.nvidia.com/nccl>`_, for use in CUDA environments.
-
-The list of all currently supported operations and their documentation can be
-seen in the :ref:`API docs<distributed>`.
+MLX utilizes `MPI <https://en.wikipedia.org/wiki/Message_Passing_Interface>`_ to
+provide distributed communication operations that allow the computational cost
+of training or inference to be shared across many physical machines. You can
+see a list of the supported operations in the :ref:`API docs<distributed>`.
 
 .. note::
-   Some operations may not be supported or not as fast as they should be.
+   A lot of operations may not be supported or not as fast as they should be.
    We are adding more and tuning the ones we have as we are figuring out the
    best way to do distributed computing on Macs using MLX.
 
 Getting Started
 ---------------
 
-A distributed program in MLX is as simple as:
+MLX already comes with the ability to "talk" to MPI if it is installed on the
+machine. The minimal distributed program in MLX is as simple as:
 
 .. code:: python
 
@@ -37,78 +30,74 @@ A distributed program in MLX is as simple as:
     print(world.rank(), x)
 
 The program above sums the array ``mx.ones(10)`` across all
-distributed processes. However, when this script is run with ``python`` only
-one process is launched and no distributed communication takes place. Namely,
-all operations in ``mx.distributed`` are noops when the distributed group has a
-size of one. This property allows us to avoid code that checks if we are in a
-distributed setting similar to the one below:
+distributed processes. If simply run with ``python``, however, only one
+process is launched and no distributed communication takes place.
 
-.. code:: python
-
-    import mlx.core as mx
-
-    x = ...
-    world = mx.distributed.init()
-    # No need for the check we can simply do x = mx.distributed.all_sum(x)
-    if world.size() > 1:
-        x = mx.distributed.all_sum(x)
-
-Running Distributed Programs
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-MLX provides ``mlx.launch`` a helper script to launch distributed programs.
-Continuing with our initial example we can run it on localhost with 4 processes using
+To launch the program in distributed mode we need to use ``mpirun`` or
+``mpiexec`` depending on the MPI installation. The simplest possible way is the
+following:
 
 .. code:: shell
 
-    $ mlx.launch -n 4 my_script.py
-    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    $ mpirun -np 2 python test.py
+    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
+    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
 
-We can also run it on some remote hosts by providing their IPs (provided that
-the script exists on all hosts and they are reachable by ssh)
+The above launches two processes on the same (local) machine and we can see
+both standard output streams. The processes send the array of 1s to each other
+and compute the sum which is printed. Launching with ``mpirun -np 4 ...`` would
+print 4 etc.
+
+Installing MPI
+---------------
+
+MPI can be installed with Homebrew, using the Anaconda package manager or
+compiled from source. Most of our testing is done using ``openmpi`` installed
+with the Anaconda package manager as follows:
 
 .. code:: shell
 
-    $ mlx.launch --hosts ip1,ip2,ip3,ip4 my_script.py
-    3 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    2 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    1 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
-    0 array([4, 4, 4, ..., 4, 4, 4], dtype=float32)
+    $ conda install openmpi
 
-Consult the dedicated :doc:`usage guide<launching_distributed>` for more
-information on using ``mlx.launch``.
+Installing with Homebrew may require specifying the location of ``libmpi.dyld``
+so that MLX can find it and load it at runtime. This can simply be achieved by
+passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun``.
 
-Selecting Backend
-^^^^^^^^^^^^^^^^^
+.. code:: shell
 
-You can select the backend you want to use when calling :func:`init` by passing
-one of ``{'any', 'ring', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all
-available backends. If they all fail then a singleton group is created.
+    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python test.py
+
+Setting up Remote Hosts
+-----------------------
+
+MPI can automatically connect to remote hosts and set up the communication over
+the network if the remote hosts can be accessed via ssh. A good checklist to
+debug connectivity issues is the following:
+
+* ``ssh hostname`` works from all machines to all machines without asking for
+  password or host confirmation
+* ``mpirun`` is accessible on all machines. You can call ``mpirun`` using its
+  full path to force all machines to use a specific path.
+* Ensure that the ``hostname`` used by MPI is the one that you have configured
+  in the ``.ssh/config`` files on all machines.
 
 .. note::
-   After a distributed backend is successfully initialized :func:`init` will
-   return **the same backend** if called without arguments or with backend set to
-   ``any``.
+  For an example hostname ``foo.bar.com`` MPI can use only ``foo`` as
+  the hostname passed to ssh if the current hostname matches ``*.bar.com``.
 
-The following examples aim to clarify the backend initialization logic in MLX:
+An easy way to pass the host names to MPI is using a host file. A host file
+looks like the following, where ``host1`` and ``host2`` should be the fully
+qualified domain names or IPs for these hosts.
 
-.. code:: python
+.. code::
 
-    # Case 1: Initialize MPI regardless if it was possible to initialize the ring backend
-    world = mx.distributed.init(backend="mpi")
-    world2 = mx.distributed.init()  # subsequent calls return the MPI backend!
+    host1 slots=1
+    host2 slots=1
 
-    # Case 2: Initialize any backend
-    world = mx.distributed.init(backend="any")  # equivalent to no arguments
-    world2 = mx.distributed.init()  # same as above
-
-    # Case 3: Initialize both backends at the same time
-    world_mpi = mx.distributed.init(backend="mpi")
-    world_ring = mx.distributed.init(backend="ring")
-    world_any = mx.distributed.init()  # same as MPI because it was initialized first!
+When using MLX, it is very likely that you want to use 1 slot per host, ie one
+process per host.  The hostfile also needs to contain the current
+host if you want to run on the local host. Passing the host file to
+``mpirun`` is simply done using the ``--hostfile`` command line argument.
 
 Training Example
 ----------------
@@ -152,13 +141,12 @@ everything else remaining the same.
     from mlx.utils import tree_map
 
     def all_reduce_grads(grads):
-        N = mx.distributed.init().size()
+        N = mx.distributed.init()
         if N == 1:
             return grads
         return tree_map(
                 lambda x: mx.distributed.all_sum(x) / N,
-            grads
-        )
+                grads)
 
     def step(model, x, y):
         loss, grads = loss_grad_fn(model, x, y)
@@ -166,181 +154,13 @@ everything else remaining the same.
         optimizer.update(model, grads)
         return loss
 
-Utilizing ``nn.average_gradients``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Tuning All Reduce
+-----------------
 
-Although the code example above works correctly; it performs one communication
-per gradient. It is significantly more efficient to aggregate several gradients
-together and perform fewer communication steps.
+We are working on improving the performance of all reduce on MLX but for now
+the two main things one can do to extract the most out of distributed training with MLX are:
 
-This is the purpose of :func:`mlx.nn.average_gradients`. The final code looks
-almost identical to the example above:
-
-.. code:: python
-
-    model = ...
-    optimizer = ...
-    dataset = ...
-
-    def step(model, x, y):
-        loss, grads = loss_grad_fn(model, x, y)
-        grads = mx.nn.average_gradients(grads)  # <---- This line was added
-        optimizer.update(model, grads)
-        return loss
-
-    for x, y in dataset:
-        loss = step(model, x, y)
-        mx.eval(loss, model.parameters())
-
-
-Getting Started with MPI
-------------------------
-
-MLX already comes with the ability to "talk" to MPI if it is installed on the
-machine. Launching distributed MLX programs that use MPI can be done with
-``mpirun`` as expected. However, in the following examples we will be using
-``mlx.launch --backend mpi`` which takes care of some nuisances such as setting
-absolute paths for the ``mpirun`` executable and the ``libmpi.dyld`` shared
-library.
-
-The simplest possible usage is the following which, assuming the minimal
-example in the beginning of this page, should result in:
-
-.. code:: shell
-
-    $ mlx.launch --backend mpi -n 2 test.py
-    1 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-    0 array([2, 2, 2, ..., 2, 2, 2], dtype=float32)
-
-The above launches two processes on the same (local) machine and we can see
-both standard output streams. The processes send the array of 1s to each other
-and compute the sum which is printed. Launching with ``mlx.launch -n 4 ...`` would
-print 4 etc.
-
-Installing MPI
-^^^^^^^^^^^^^^
-
-MPI can be installed with Homebrew, pip, using the Anaconda package manager, or
-compiled from source. Most of our testing is done using ``openmpi`` installed
-with the Anaconda package manager as follows:
-
-.. code:: shell
-
-    $ conda install conda-forge::openmpi
-
-Installing with Homebrew or pip requires specifying the location of ``libmpi.dyld``
-so that MLX can find it and load it at runtime. This can simply be achieved by
-passing the ``DYLD_LIBRARY_PATH`` environment variable to ``mpirun`` and it is
-done automatically by ``mlx.launch``. Some environments use a non-standard
-library filename that can be specified using the ``MPI_LIBNAME`` environment
-variable. This is automatically taken care of by ``mlx.launch`` as well.
-
-.. code:: shell
-
-    $ mpirun -np 2 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ -x MPI_LIBNAME=libmpi.40.dylib python test.py
-    $ # or simply
-    $ mlx.launch -n 2 test.py
-
-Setting up Remote Hosts
-^^^^^^^^^^^^^^^^^^^^^^^
-
-MPI can automatically connect to remote hosts and set up the communication over
-the network if the remote hosts can be accessed via ssh. A good checklist to
-debug connectivity issues is the following:
-
-* ``ssh hostname`` works from all machines to all machines without asking for
-  password or host confirmation
-* ``mpirun`` is accessible on all machines.
-* Ensure that the ``hostname`` used by MPI is the one that you have configured
-  in the ``.ssh/config`` files on all machines.
-
-Tuning MPI All Reduce
-^^^^^^^^^^^^^^^^^^^^^
-
-.. note::
-
-    For faster all reduce consider using the ring backend either with Thunderbolt
-    connections or over Ethernet.
-
-Configure MPI to use N tcp connections between each host to improve bandwidth
-by passing ``--mca btl_tcp_links N``.
-
-Force MPI to use the most performant network interface by setting ``--mca
-btl_tcp_if_include <iface>`` where ``<iface>`` should be the interface you want
-to use.
-
-Getting Started with Ring
--------------------------
-
-The ring backend does not depend on any third party library so it is always
-available. It uses TCP sockets so the nodes need to be reachable via a network.
-As the name suggests the nodes are connected in a ring which means that rank 1
-can only communicate with rank 0 and rank 2, rank 2 only with rank 1 and rank 3
-and so on and so forth. As a result :func:`send` and :func:`recv` with
-arbitrary sender and receiver is not supported in the ring backend.
-
-Defining a Ring
-^^^^^^^^^^^^^^^
-
-The easiest way to define and use a ring is via a JSON hostfile and the
-``mlx.launch`` :doc:`helper script <launching_distributed>`. For each node one
-defines a hostname to ssh into to run commands on this node and one or more IPs
-that this node will listen to for connections.
-
-For example the hostfile below defines a 4 node ring. ``hostname1`` will be
-rank 0, ``hostname2`` rank 1 etc.
-
-.. code:: json
-
-    [
-        {"ssh": "hostname1", "ips": ["123.123.123.1"]},
-        {"ssh": "hostname2", "ips": ["123.123.123.2"]},
-        {"ssh": "hostname3", "ips": ["123.123.123.3"]},
-        {"ssh": "hostname4", "ips": ["123.123.123.4"]}
-    ]
-
-Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each
-node, run the script which will listen for connections in each of the provided
-IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a
-connection from ``123.123.123.4`` and so on and so forth.
-
-Thunderbolt Ring
-^^^^^^^^^^^^^^^^
-
-Although the ring backend can have benefits over MPI even for Ethernet, its
-main purpose is to use Thunderbolt rings for higher bandwidth communication.
-Setting up such thunderbolt rings can be done manually, but is a relatively
-tedious process. To simplify this, we provide the utility ``mlx.distributed_config``.
-
-To use ``mlx.distributed_config`` your computers need to be accessible by ssh via
-Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the
-utility as follows:
-
-.. code:: shell
-
-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4
-
-By default the script will attempt to discover the thunderbolt ring and provide
-you with the commands to configure each node as well as the ``hostfile.json``
-to use with ``mlx.launch``. If password-less ``sudo`` is available on the nodes
-then ``--auto-setup`` can be used to configure them automatically.
-
-To validate your connection without configuring anything
-``mlx.distributed_config`` can also plot the ring using DOT format.
-
-.. code:: shell
-
-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot
-   dot -Tpng ring.dot >ring.png
-   open ring.png
-
-If you want to go through the process manually, the steps are as follows:
-
-* Disable the thunderbolt bridge interface
-* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces
-  corresponding to that cable in nodes ``i`` and ``i + 1``.
-* Set up a unique subnetwork connecting the two nodes for the corresponding
-  interfaces. For instance if the cable corresponds to ``en2`` on node ``i``
-  and ``en2`` also on node ``i + 1`` then we may assign IPs ``192.168.0.1`` and
-  ``192.168.0.2`` respectively to the two nodes. For more details you can see
-  the commands prepared by the utility script.
+1. Perform a few large reductions instead of many small ones to improve
+   bandwidth and latency
+2. Pass ``--mca btl_tcp_links 4`` to ``mpirun`` to configure it to use 4 tcp
+   connections between each host to improve bandwidth

docs/src/usage/export.rst

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

docs/src/usage/indexing.rst

@@ -70,8 +70,7 @@ Differences from NumPy
 
   * Indexing does not perform bounds checking. Indexing out of bounds is
     undefined behavior.
-  * Boolean mask based indexing is supported for assignment only (see
-    :ref:`boolean-mask-assignment`).
+  * Boolean mask based indexing is not yet supported.
 
 The reason for the lack of bounds checking is that exceptions cannot propagate
 from the GPU. Performing bounds checking for array indices before launching the
@@ -108,28 +107,6 @@ same array:
   >>> a
   array([1, 2, 0], dtype=int32)
 
-Note that unlike NumPy, slicing an array creates a copy, not a view. So
-mutating it does not mutate the original array:
-
-.. code-block:: shell
-
-  >>> a = mx.array([1, 2, 3])
-  >>> b = a[:]
-  >>> b[2] = 0
-  >>> b
-  array([1, 2, 0], dtype=int32)
-  >>> a
-  array([1, 2, 3], dtype=int32)
-
-Also unlike NumPy, updates to the same location are nondeterministic:
-
-.. code-block:: shell
-
-  >>> a = mx.array([1, 2, 3])
-  >>> a[[0, 0]] = mx.array([4, 5])
-
-The first element of ``a`` could be ``4`` or ``5``.
-
 Transformations of functions which use in-place updates are allowed and work as
 expected. For example:
 
@@ -144,51 +121,3 @@ expected. For example:
 
 In the above ``dfdx`` will have the correct gradient, namely zeros at ``idx``
 and ones elsewhere.
-
-.. _boolean-mask-assignment:
-
-Boolean Mask Assignment
------------------------
-
-MLX supports boolean indices using NumPy syntax. A mask must already be
-a :class:`bool_` MLX :class:`array` or a NumPy ``ndarray`` with ``dtype=bool``.
-Other index types are routed through the standard scatter code.
-
-.. code-block:: shell
-
-   >>> a = mx.array([1.0, 2.0, 3.0])
-   >>> mask = mx.array([True, False, True])
-   >>> updates = mx.array([5.0, 6.0])
-   >>> a[mask] = updates
-   >>> a
-   array([5.0, 2.0, 6.0], dtype=float32)
-
-Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar
-assignments, ``updates`` must provide at least as many elements as there are
-``True`` entries in ``mask``.
-
-.. code-block:: shell
-
-   >>> a = mx.zeros((2, 3))
-   >>> mask = mx.array([[True, False, True],
-                        [False, False, True]])
-   >>> a[mask] = 1.0
-   >>> a
-   array([[1.0, 0.0, 1.0],
-          [0.0, 0.0, 1.0]], dtype=float32)
-
-Boolean masks follow NumPy semantics:
-
-- The mask shape must match the shape of the axes it indexes exactly. No mask
-  broadcasting occurs.
-- Any axes not covered by the mask are taken in full.
-
-.. code-block:: shell
-
-   >>> a = mx.arange(1000).reshape(10, 10, 10)
-   >>> a[mx.random.randn(10, 10) > 0.0] = 0  # valid: mask covers axes 0 and 1
-
-The mask of shape ``(10, 10)`` applies to the first two axes, so ``a[mask]``
-selects the 1-D slices ``a[i, j, :]`` where ``mask[i, j]`` is ``True``.
-Shapes such as ``(1, 10, 10)`` or ``(10, 10, 1)`` do not match the indexed
-axes and therefore raise errors.

docs/src/usage/launching_distributed.rst

@@ -1,105 +0,0 @@
-:orphan:
-
-.. _usage_launch_distributed:
-
-Launching Distributed Programs
-==============================
-
-.. currentmodule:: mlx.core.distributed
-
-Installing the MLX python package provides a helper script ``mlx.launch`` that
-can be used to run python scripts distributed on several nodes. It allows
-launching using either the MPI backend or the ring backend. See the
-:doc:`distributed docs <distributed>` for the different backends.
-
-Usage
------
-
-The minimal usage example of ``mlx.launch`` is simply
-
-.. code:: shell
-
-    mlx.launch --hosts ip1,ip2 my_script.py
-
-or for testing on localhost
-
-.. code:: shell
-
-    mlx.launch -n 2 my_script.py
-
-The ``mlx.launch`` command connects to the provided host and launches the input
-script on each host. It monitors each of the launched processes and terminates
-the rest if one of them fails unexpectedly or if ``mlx.launch`` is terminated.
-It also takes care of forwarding the output of each remote process to stdout
-and stderr respectively.
-
-Providing Hosts
-^^^^^^^^^^^^^^^^
-
-Hosts can be provided as command line arguments, like above, but the way that
-allows to fully define a list of hosts is via a JSON hostfile. The hostfile has
-a very simple schema. It is simply a list of objects that define each host via
-a hostname to ssh to and a list of IPs to utilize for the communication.
-
-.. code:: json
-
-    [
-        {"ssh": "hostname1", "ips": ["123.123.1.1", "123.123.2.1"]},
-        {"ssh": "hostname2", "ips": ["123.123.1.2", "123.123.2.2"]}
-    ]
-
-You can use ``mlx.distributed_config --over ethernet`` to create a hostfile
-with IPs corresponding to the ``en0`` interface.
-
-Setting up Remote Hosts
-^^^^^^^^^^^^^^^^^^^^^^^^
-
-In order to be able to launch the script on each host we need to be able to
-connect via ssh. Moreover the input script and python binary need to be on each
-host and on the same path. A good checklist to debug errors is the following:
-
-* ``ssh hostname`` works without asking for password or host confirmation
-* the python binary is available on all hosts at the same path. You can use
-  ``mlx.launch --print-python`` to see what that path is.
-* the script you want to run is available on all hosts at the same path
-
-.. _mpi_specifics:
-
-MPI Specifics
--------------
-
-One can use MPI by passing ``--backend mpi`` to ``mlx.launch``. In that case,
-``mlx.launch`` is a thin wrapper over ``mpirun``. Moreover,
-
-* The IPs in the hostfile are ignored
-* The ssh connectivity requirement is stronger as every node needs to be able
-  to connect to every other node
-* ``mpirun`` needs to be available on every node at the same path
-
-Finally, one can pass arguments to ``mpirun`` using ``--mpi-arg``. For instance
-to choose a specific interface for the byte-transfer-layer of MPI we can call
-``mlx.launch`` as follows:
-
-.. code:: shell
-
-    mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py
-
-
-.. _ring_specifics:
-
-Ring Specifics
---------------
-
-The ring backend, which is also the default backend, can be explicitly selected
-with the argument ``--backend ring``. The ring backend has some specific
-requirements and arguments that are different to MPI:
-
-* The argument ``--hosts`` only accepts IPs and not hostnames. If we need to
-  ssh to a hostname that does not correspond to the IP we want to bind to we
-  have to provide a hostfile.
-* ``--starting-port`` defines the port to bind to on the remote hosts.
-  Specifically rank 0 for the first IP will use this port and each subsequent
-  IP or rank will add 1 to this port.
-* ``--connections-per-ip`` allows us to increase the number of connections
-  between neighboring nodes. This corresponds to ``--mca btl_tcp_links 2`` for
-  ``mpirun``.

docs/src/usage/numpy.rst

@@ -21,13 +21,11 @@ Let's convert an array to NumPy and back.
 
 .. note::
 
-    Since NumPy does not support ``bfloat16`` arrays, you will need to convert
-    to ``float16`` or ``float32`` first: ``np.array(a.astype(mx.float32))``.
-    Otherwise, you will receive an error like: ``Item size 2 for PEP 3118
-    buffer format string does not match the dtype V item size 0.``
+    Since NumPy does not support ``bfloat16`` arrays, you will need to convert to ``float16`` or ``float32`` first:
+    ``np.array(a.astype(mx.float32))``.
+    Otherwise, you will receive an error like: ``Item size 2 for PEP 3118 buffer format string does not match the dtype V item size 0.``
 
-By default, NumPy copies data to a new array. This can be prevented by creating
-an array view:
+By default, NumPy copies data to a new array. This can be prevented by creating an array view:
 
 .. code-block:: python
 
@@ -37,16 +35,10 @@ an array view:
   a_view[0] = 1
   print(a[0].item()) # 1
 
-.. note::
+A NumPy array view is a normal NumPy array, except that it does not own its memory.
+This means writing to the view is reflected in the original array.
 
-    NumPy arrays with type ``float64`` will be default converted to MLX arrays
-    with type ``float32``.
-
-A NumPy array view is a normal NumPy array, except that it does not own its
-memory. This means writing to the view is reflected in the original array.
-
-While this is quite powerful to prevent copying arrays, it should be noted that
-external changes to the memory of arrays cannot be reflected in gradients.
+While this is quite powerful to prevent copying arrays, it should be noted that external changes to the memory of arrays cannot be reflected in gradients.
 
 Let's demonstrate this in an example:
 
@@ -64,12 +56,11 @@ Let's demonstrate this in an example:
 
 
 The function ``f`` indirectly modifies the array ``x`` through a memory view.
-However, this modification is not reflected in the gradient, as seen in the
-last line outputting ``1.0``, representing the gradient of the sum operation
-alone.  The squaring of ``x`` occurs externally to MLX, meaning that no
-gradient is incorporated.  It's important to note that a similar issue arises
-during array conversion and copying.  For instance, a function defined as
-``mx.array(np.array(x)**2).sum()`` would also result in an incorrect gradient,
+However, this modification is not reflected in the gradient, as seen in the last line outputting ``1.0``,
+representing the gradient of the sum operation alone.
+The squaring of ``x`` occurs externally to MLX, meaning that no gradient is incorporated.
+It's important to note that a similar issue arises during array conversion and copying.
+For instance, a function defined as ``mx.array(np.array(x)**2).sum()`` would also result in an incorrect gradient,
 even though no in-place operations on MLX memory are executed.
 
 PyTorch
@@ -80,8 +71,7 @@ PyTorch
    PyTorch Support for :obj:`memoryview` is experimental and can break for
    multi-dimensional arrays. Casting to NumPy first is advised for now.
 
-PyTorch supports the buffer protocol, but it requires an explicit
-:obj:`memoryview`.
+PyTorch supports the buffer protocol, but it requires an explicit :obj:`memoryview`.
 
 .. code-block:: python
 
@@ -92,8 +82,7 @@ PyTorch supports the buffer protocol, but it requires an explicit
   b = torch.tensor(memoryview(a))
   c = mx.array(b.numpy())
 
-Conversion from PyTorch tensors back to arrays must be done via intermediate
-NumPy arrays with ``numpy()``.
+Conversion from PyTorch tensors back to arrays must be done via intermediate NumPy arrays with ``numpy()``.
 
 JAX
 ---
@@ -111,8 +100,7 @@ JAX fully supports the buffer protocol.
 TensorFlow
 ----------
 
-TensorFlow supports the buffer protocol, but it requires an explicit
-:obj:`memoryview`.
+TensorFlow supports the buffer protocol, but it requires an explicit :obj:`memoryview`.
 
 .. code-block:: python
 

examples/cmake_project/CMakeLists.txt

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

examples/cmake_project/README.md

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

examples/cmake_project/example.cpp

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

examples/cpp/distributed.cpp

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

examples/cpp/linear_regression.cpp

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

examples/cpp/logistic_regression.cpp

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

examples/cpp/metal_capture.cpp

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

examples/cpp/tutorial.cpp

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

examples/export/CMakeLists.txt

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

examples/export/README.md

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

examples/export/eval_mlp.cpp

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

examples/export/eval_mlp.py

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

examples/export/train_mlp.cpp

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

examples/export/train_mlp.py

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

examples/extensions/CMakeLists.txt

@@ -10,6 +10,7 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
 
 # ----------------------------- Dependencies -----------------------------
+find_package(MLX CONFIG REQUIRED)
 find_package(
   Python 3.8
   COMPONENTS Interpreter Development.Module
@@ -17,15 +18,10 @@ find_package(
 execute_process(
   COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
   OUTPUT_STRIP_TRAILING_WHITESPACE
-  OUTPUT_VARIABLE nanobind_ROOT)
+  OUTPUT_VARIABLE NB_DIR)
+list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
 find_package(nanobind CONFIG REQUIRED)
 
-execute_process(
-  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
-  OUTPUT_STRIP_TRAILING_WHITESPACE
-  OUTPUT_VARIABLE MLX_ROOT)
-find_package(MLX CONFIG REQUIRED)
-
 # ----------------------------- Extensions -----------------------------
 
 # Add library

examples/extensions/axpby/axpby.cpp

@@ -1,34 +1,25 @@
-// Copyright © 2023-2025 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
-#include <dlfcn.h>
+#include <cassert>
 #include <iostream>
 #include <sstream>
 
+#include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/utils.h"
-#include "mlx/backend/cpu/encoder.h"
 #include "mlx/utils.h"
 
 #include "axpby/axpby.h"
 
+#ifdef ACCELERATE_NEW_LAPACK
+#include <vecLib/cblas_new.h>
+#endif
+
 #ifdef _METAL_
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/utils.h"
 #endif
 
-namespace my_ext {
-
-// A helper function to find the location of the current binary on disk.
-// The Metal library ("mlx_ext.mtllib"), should be in the same directory.
-std::string current_binary_dir() {
-  static std::string binary_dir = []() {
-    Dl_info info;
-    if (!dladdr(reinterpret_cast<void*>(&current_binary_dir), &info)) {
-      throw std::runtime_error("Unable to get current binary dir.");
-    }
-    return std::filesystem::path(info.dli_fname).parent_path().string();
-  }();
-  return binary_dir;
-}
+namespace mlx::core {
 
 ///////////////////////////////////////////////////////////////////////////////
 // Operation Implementation
@@ -41,24 +32,24 @@ std::string current_binary_dir() {
  *  Follow numpy style broadcasting between x and y
  *  Inputs are upcasted to floats if needed
  **/
-mx::array axpby(
-    const mx::array& x, // Input mx::array x
-    const mx::array& y, // Input mx::array y
+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
-    mx::StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
+    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 = mx::issubdtype(promoted_dtype, mx::float32)
+  auto out_dtype = issubdtype(promoted_dtype, float32)
       ? promoted_dtype
-      : promote_types(promoted_dtype, mx::float32);
+      : promote_types(promoted_dtype, float32);
 
   // Cast x and y up to the determined dtype (on the same stream s)
-  auto x_casted = mx::astype(x, out_dtype, s);
-  auto y_casted = mx::astype(y, out_dtype, 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);
@@ -66,12 +57,12 @@ mx::array axpby(
 
   // Construct the array as the output of the Axpby primitive
   // with the broadcasted and upcasted arrays as inputs
-  return mx::array(
-      /* const mx::Shape& shape = */ out_shape,
-      /* mx::Dtype dtype = */ out_dtype,
-      /* std::shared_ptr<mx::Primitive> primitive = */
+  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<mx::array>& inputs = */ broadcasted_inputs);
+      /* const std::vector<array>& inputs = */ broadcasted_inputs);
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -80,69 +71,140 @@ mx::array axpby(
 
 template <typename T>
 void axpby_impl(
-    const mx::array& x,
-    const mx::array& y,
-    mx::array& out,
+    const array& x,
+    const array& y,
+    array& out,
     float alpha_,
-    float beta_,
-    mx::Stream stream) {
-  out.set_data(mx::allocator::malloc(out.nbytes()));
+    float beta_) {
+  // We only allocate memory when we are ready to fill the output
+  // malloc_or_wait synchronously allocates available memory
+  // There may be a wait executed here if the allocation is requested
+  // under memory-pressured conditions
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
-  // Get the CPU command encoder and register input and output arrays
-  auto& encoder = mx::cpu::get_command_encoder(stream);
-  encoder.set_input_array(x);
-  encoder.set_input_array(y);
-  encoder.set_output_array(out);
+  // Collect input and output data pointers
+  const T* x_ptr = x.data<T>();
+  const T* y_ptr = y.data<T>();
+  T* out_ptr = out.data<T>();
 
-  // Launch the CPU kernel
-  encoder.dispatch([x_ptr = x.data<T>(),
-                    y_ptr = y.data<T>(),
-                    out_ptr = out.data<T>(),
-                    size = out.size(),
-                    shape = out.shape(),
-                    x_strides = x.strides(),
-                    y_strides = y.strides(),
-                    alpha_,
-                    beta_]() {
   // Cast alpha and beta to the relevant types
   T alpha = static_cast<T>(alpha_);
   T beta = static_cast<T>(beta_);
 
   // Do the element-wise operation for each output
-    for (size_t out_idx = 0; out_idx < size; out_idx++) {
+  for (size_t out_idx = 0; out_idx < out.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);
+    auto x_offset = elem_to_loc(out_idx, x.shape(), x.strides());
+    auto y_offset = elem_to_loc(out_idx, y.shape(), y.strides());
 
     // We allocate the output to be contiguous and regularly strided
     // (defaults to row major) and hence it doesn't need additional mapping
     out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
   }
-  });
 }
 
-void Axpby::eval_cpu(
-    const std::vector<mx::array>& inputs,
-    std::vector<mx::array>& outputs) {
+/** Fall back implementation for evaluation on CPU */
+void Axpby::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  // Check the inputs (registered in the op while constructing the out array)
+  assert(inputs.size() == 2);
   auto& x = inputs[0];
   auto& y = inputs[1];
   auto& out = outputs[0];
 
   // Dispatch to the correct dtype
-  if (out.dtype() == 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());
+  if (out.dtype() == float32) {
+    return axpby_impl<float>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == float16) {
+    return axpby_impl<float16_t>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == bfloat16) {
+    return axpby_impl<bfloat16_t>(x, y, out, alpha_, beta_);
+  } else if (out.dtype() == complex64) {
+    return axpby_impl<complex64_t>(x, y, out, alpha_, beta_);
   } else {
     throw std::runtime_error(
         "Axpby is only supported for floating point types.");
   }
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// Primitive Accelerate Backend Implementation
+///////////////////////////////////////////////////////////////////////////////
+
+#ifdef ACCELERATE_NEW_LAPACK
+
+template <typename T>
+void axpby_impl_accelerate(
+    const array& x,
+    const array& y,
+    array& out,
+    float alpha_,
+    float beta_) {
+  // Accelerate library provides catlas_saxpby which does
+  // Y = (alpha * X) + (beta * Y) in place
+  // To use it, we first copy the data in y over to the output array
+
+  // This specialization requires both x and y be contiguous in the same mode
+  // i.e: corresponding linear indices in both point to corresponding elements
+  // The data in the output array is allocated to match the strides in y
+  // such that x, y, and out are contiguous in the same mode and
+  // no transposition is needed
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+
+  // We then copy over the elements using the contiguous vector specialization
+  copy_inplace(y, out, CopyType::Vector);
+
+  // Get x and y pointers for catlas_saxpby
+  const T* x_ptr = x.data<T>();
+  T* y_ptr = out.data<T>();
+
+  T alpha = static_cast<T>(alpha_);
+  T beta = static_cast<T>(beta_);
+
+  // Call the inplace accelerate operator
+  catlas_saxpby(
+      /* N = */ out.size(),
+      /* ALPHA = */ alpha,
+      /* X = */ x_ptr,
+      /* INCX = */ 1,
+      /* BETA = */ beta,
+      /* Y = */ y_ptr,
+      /* INCY = */ 1);
+}
+
+/** Evaluate primitive on CPU using accelerate specializations */
+void Axpby::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 2);
+  auto& x = inputs[0];
+  auto& y = inputs[1];
+  auto& out = outputs[0];
+
+  // Accelerate specialization for contiguous single precision float arrays
+  if (out.dtype() == float32 &&
+      ((x.flags().row_contiguous && y.flags().row_contiguous) ||
+       (x.flags().col_contiguous && y.flags().col_contiguous))) {
+    axpby_impl_accelerate<float>(x, y, out, alpha_, beta_);
+    return;
+  }
+
+  // Fall back to common backend if specializations are not available
+  eval(inputs, outputs);
+}
+
+#else // Accelerate not available
+
+/** Evaluate primitive on CPU falling back to common backend */
+void Axpby::eval_cpu(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+
+#endif
+
 ///////////////////////////////////////////////////////////////////////////////
 // Primitive Metal Backend Implementation
 ///////////////////////////////////////////////////////////////////////////////
@@ -151,9 +213,10 @@ void Axpby::eval_cpu(
 
 /** Evaluate primitive on GPU */
 void Axpby::eval_gpu(
-    const std::vector<mx::array>& inputs,
-    std::vector<mx::array>& outputs) {
+    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];
@@ -162,7 +225,7 @@ void Axpby::eval_gpu(
   // and each stream carries its device identifiers
   auto& s = stream();
   // We get the needed metal device using the stream
-  auto& d = mx::metal::device(s.device);
+  auto& d = metal::device(s.device);
 
   // Prepare to specialize based on contiguity
   bool contiguous_kernel =
@@ -172,24 +235,25 @@ void Axpby::eval_gpu(
   // Allocate output memory with strides based on specialization
   if (contiguous_kernel) {
     out.set_data(
-        mx::allocator::malloc(x.data_size() * out.itemsize()),
+        allocator::malloc_or_wait(x.data_size() * out.itemsize()),
         x.data_size(),
         x.strides(),
         x.flags());
   } else {
-    out.set_data(mx::allocator::malloc(out.nbytes()));
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
 
   // Resolve name of kernel (corresponds to axpby.metal)
-  std::string kname = "axpby_";
-  kname += (contiguous_kernel ? "contiguous_" : "general_");
-  kname += type_to_name(out);
+  std::ostringstream kname;
+  kname << "axpby_";
+  kname << (contiguous_kernel ? "contiguous_" : "general_");
+  kname << type_to_name(out);
 
-  // Load the metal library
-  auto lib = d.get_library("mlx_ext", current_binary_dir());
+  // Make sure the metal library is available
+  d.register_library("mlx_ext");
 
   // Make a kernel from this metal library
-  auto kernel = d.get_kernel(kname, lib);
+  auto kernel = d.get_kernel(kname.str(), "mlx_ext");
 
   // Prepare to encode kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
@@ -215,7 +279,7 @@ void Axpby::eval_gpu(
   if (!contiguous_kernel) {
     compute_encoder.set_vector_bytes(x.shape(), 5);
     compute_encoder.set_vector_bytes(x.strides(), 6);
-    compute_encoder.set_vector_bytes(y.strides(), 7);
+    compute_encoder.set_bytes(y.strides(), 7);
     compute_encoder.set_bytes(ndim, 8);
   }
 
@@ -238,8 +302,8 @@ void Axpby::eval_gpu(
 
 /** Fail evaluation on GPU */
 void Axpby::eval_gpu(
-    const std::vector<mx::array>& inputs,
-    std::vector<mx::array>& out) {
+    const std::vector<array>& inputs,
+    std::vector<array>& out) {
   throw std::runtime_error("Axpby has no GPU implementation.");
 }
 
@@ -250,9 +314,9 @@ void Axpby::eval_gpu(
 ///////////////////////////////////////////////////////////////////////////////
 
 /** The Jacobian-vector product. */
-std::vector<mx::array> Axpby::jvp(
-    const std::vector<mx::array>& primals,
-    const std::vector<mx::array>& tangents,
+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 built with ops
@@ -264,8 +328,8 @@ std::vector<mx::array> Axpby::jvp(
   // scaled by beta
   if (argnums.size() > 1) {
     auto scale = argnums[0] == 0 ? alpha_ : beta_;
-    auto scale_arr = mx::array(scale, tangents[0].dtype());
-    return {mx::multiply(scale_arr, tangents[0], stream())};
+    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
@@ -275,24 +339,24 @@ std::vector<mx::array> Axpby::jvp(
 }
 
 /** The vector-Jacobian product. */
-std::vector<mx::array> Axpby::vjp(
-    const std::vector<mx::array>& primals,
-    const std::vector<mx::array>& cotangents,
+std::vector<array> Axpby::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
     const std::vector<int>& argnums,
-    const std::vector<mx::array>&) {
+    const std::vector<array>&) {
   // Reverse mode diff
-  std::vector<mx::array> vjps;
+  std::vector<array> vjps;
   for (auto arg : argnums) {
     auto scale = arg == 0 ? alpha_ : beta_;
-    auto scale_arr = mx::array(scale, cotangents[0].dtype());
-    vjps.push_back(mx::multiply(scale_arr, cotangents[0], stream()));
+    auto scale_arr = array(scale, cotangents[0].dtype());
+    vjps.push_back(multiply(scale_arr, cotangents[0], stream()));
   }
   return vjps;
 }
 
 /** Vectorize primitive along given axis */
-std::pair<std::vector<mx::array>, std::vector<int>> Axpby::vmap(
-    const std::vector<mx::array>& inputs,
+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 has no vmap implementation.");
 }
@@ -303,4 +367,4 @@ bool Axpby::is_equivalent(const Primitive& other) const {
   return alpha_ == r_other.alpha_ && beta_ == r_other.beta_;
 }
 
-} // namespace my_ext
+} // namespace mlx::core

examples/extensions/axpby/axpby.h

@@ -1,13 +1,11 @@
-// Copyright © 2023-2025 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #pragma once
 
 #include "mlx/ops.h"
 #include "mlx/primitives.h"
 
-namespace mx = mlx::core;
-
-namespace my_ext {
+namespace mlx::core {
 
 ///////////////////////////////////////////////////////////////////////////////
 // Operation
@@ -20,22 +18,22 @@ namespace my_ext {
  *  Follow numpy style broadcasting between x and y
  *  Inputs are upcasted to floats if needed
  **/
-mx::array axpby(
-    const mx::array& x, // Input array x
-    const mx::array& y, // Input array y
+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
-    mx::StreamOrDevice s = {} // Stream on which to schedule the operation
+    StreamOrDevice s = {} // Stream on which to schedule the operation
 );
 
 ///////////////////////////////////////////////////////////////////////////////
 // Primitive
 ///////////////////////////////////////////////////////////////////////////////
 
-class Axpby : public mx::Primitive {
+class Axpby : public Primitive {
  public:
-  explicit Axpby(mx::Stream stream, float alpha, float beta)
-      : mx::Primitive(stream), alpha_(alpha), beta_(beta) {};
+  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
@@ -44,25 +42,23 @@ class Axpby : public mx::Primitive {
    * To avoid unnecessary allocations, the evaluation function
    * is responsible for allocating space for the array.
    */
-  void eval_cpu(
-      const std::vector<mx::array>& inputs,
-      std::vector<mx::array>& outputs) override;
-  void eval_gpu(
-      const std::vector<mx::array>& inputs,
-      std::vector<mx::array>& outputs) override;
+  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<mx::array> jvp(
-      const std::vector<mx::array>& primals,
-      const std::vector<mx::array>& tangents,
+  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<mx::array> vjp(
-      const std::vector<mx::array>& primals,
-      const std::vector<mx::array>& cotangents,
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
       const std::vector<int>& argnums,
-      const std::vector<mx::array>& outputs) override;
+      const std::vector<array>& outputs) override;
 
   /**
    * The primitive must know how to vectorize itself across
@@ -70,21 +66,24 @@ class Axpby : public mx::Primitive {
    * representing the vectorized computation and the axis which
    * corresponds to the output vectorized dimension.
    */
-  std::pair<std::vector<mx::array>, std::vector<int>> vmap(
-      const std::vector<mx::array>& inputs,
+  std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
       const std::vector<int>& axes) override;
 
-  /** The name of primitive. */
-  const char* name() const override {
-    return "Axpby";
+  /** Print the primitive. */
+  void print(std::ostream& os) override {
+    os << "Axpby";
   }
 
   /** Equivalence check **/
-  bool is_equivalent(const mx::Primitive& other) const override;
+  bool is_equivalent(const Primitive& other) const override;
 
  private:
   float alpha_;
   float beta_;
+
+  /** Fall back implementation for evaluation on CPU */
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
 };
 
-} // namespace my_ext
+} // namespace mlx::core

examples/extensions/axpby/axpby.metal

@@ -1,4 +1,4 @@
-// Copyright © 2023-2025 Apple Inc.
+// Copyright © 2023 Apple Inc.
 
 #include <metal_stdlib>
 
@@ -12,8 +12,8 @@ template <typename T>
     constant const float& alpha [[buffer(3)]],
     constant const float& beta [[buffer(4)]],
     constant const int* shape [[buffer(5)]],
-    constant const int64_t* x_strides [[buffer(6)]],
-    constant const int64_t* y_strides [[buffer(7)]],
+    constant const size_t* x_strides [[buffer(6)]],
+    constant const size_t* y_strides [[buffer(7)]],
     constant const int& ndim [[buffer(8)]],
     uint index [[thread_position_in_grid]]) {
   auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
@@ -34,14 +34,29 @@ template <typename T>
       static_cast<T>(alpha) * x[index] + static_cast<T>(beta) * y[index];
 }
 
-// clang-format off
 #define instantiate_axpby(type_name, type)                               \
-  instantiate_kernel("axpby_general_" #type_name, axpby_general, type) \
-  instantiate_kernel(                                                  \
-          "axpby_contiguous_" #type_name, axpby_contiguous, type)
+  template [[host_name("axpby_general_" #type_name)]] [[kernel]] void    \
+  axpby_general<type>(                                                   \
+      device const type* x [[buffer(0)]],                                \
+      device const type* y [[buffer(1)]],                                \
+      device type* out [[buffer(2)]],                                    \
+      constant const float& alpha [[buffer(3)]],                         \
+      constant const float& beta [[buffer(4)]],                          \
+      constant const int* shape [[buffer(5)]],                           \
+      constant const size_t* x_strides [[buffer(6)]],                    \
+      constant const size_t* y_strides [[buffer(7)]],                    \
+      constant const int& ndim [[buffer(8)]],                            \
+      uint index [[thread_position_in_grid]]);                           \
+  template [[host_name("axpby_contiguous_" #type_name)]] [[kernel]] void \
+  axpby_contiguous<type>(                                                \
+      device const type* x [[buffer(0)]],                                \
+      device const type* y [[buffer(1)]],                                \
+      device type* out [[buffer(2)]],                                    \
+      constant const float& alpha [[buffer(3)]],                         \
+      constant const float& beta [[buffer(4)]],                          \
+      uint index [[thread_position_in_grid]]);
 
 instantiate_axpby(float32, float);
 instantiate_axpby(float16, half);
 instantiate_axpby(bfloat16, bfloat16_t);
 instantiate_axpby(complex64, complex64_t);
-// clang-format on

examples/extensions/bindings.cpp

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

examples/extensions/pyproject.toml

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

examples/extensions/requirements.txt

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

examples/extensions/test.py

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

mlx/CMakeLists.txt

@@ -5,8 +5,6 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/dtype_utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/export.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
@@ -20,48 +18,24 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
 
-# Define MLX_VERSION only in the version.cpp file.
-add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp)
-target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}")
-target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:mlx_version>)
-
-if(MSVC)
-  # Disable some MSVC warnings to speed up compilation.
-  target_compile_options(mlx PUBLIC /wd4068 /wd4244 /wd4267 /wd4804)
-endif()
-
-if(WIN32)
-  # Export symbols by default to behave like macOS/linux.
-  set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE)
-endif()
-
-add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
-
 if(MLX_BUILD_CPU)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cpu)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
 else()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_cpu)
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
+if(MLX_BUILD_ACCELERATE)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/accelerate)
+elseif(MLX_BUILD_CPU)
+  target_sources(
+    mlx
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/common/default_primitives.cpp)
+endif()
 
 if(MLX_BUILD_METAL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
 else()
-  target_sources(mlx
-                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/no_metal.cpp)
-endif()
-
-if(MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda)
-else()
-  target_sources(mlx
-                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda/no_cuda.cpp)
-endif()
-
-if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
-else()
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_gpu)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_metal)
 endif()

mlx/allocator.cpp

@@ -4,11 +4,12 @@
 #include <sstream>
 
 #include "mlx/allocator.h"
+#include "mlx/scheduler.h"
 
 namespace mlx::core::allocator {
 
 Buffer malloc(size_t size) {
-  auto buffer = allocator().malloc(size);
+  auto buffer = allocator().malloc(size, /* allow_swap */ true);
   if (size && !buffer.ptr()) {
     std::ostringstream msg;
     msg << "[malloc] Unable to allocate " << size << " bytes.";
@@ -21,4 +22,45 @@ void free(Buffer buffer) {
   allocator().free(buffer);
 }
 
+Buffer CommonAllocator::malloc(size_t size, bool) {
+  void* ptr = std::malloc(size + sizeof(size_t));
+  if (ptr != nullptr) {
+    *static_cast<size_t*>(ptr) = size;
+  }
+  return Buffer{ptr};
+}
+
+void CommonAllocator::free(Buffer buffer) {
+  std::free(buffer.ptr());
+}
+
+size_t CommonAllocator::size(Buffer buffer) const {
+  if (buffer.ptr() == nullptr) {
+    return 0;
+  }
+  return *static_cast<size_t*>(buffer.ptr());
+}
+
+Buffer malloc_or_wait(size_t size) {
+  auto buffer = allocator().malloc(size);
+
+  while (size && !buffer.ptr() && scheduler::n_active_tasks() > 0) {
+    scheduler::wait_for_one();
+    buffer = allocator().malloc(size);
+  }
+
+  // Try swapping if needed
+  if (size && !buffer.ptr()) {
+    buffer = allocator().malloc(size, /* allow_swap = */ true);
+  }
+
+  if (size && !buffer.ptr()) {
+    std::ostringstream msg;
+    msg << "[malloc_or_wait] Unable to allocate " << size << " bytes.";
+    throw std::runtime_error(msg.str());
+  }
+
+  return buffer;
+}
+
 } // namespace mlx::core::allocator

mlx/allocator.h

@@ -14,7 +14,7 @@ class Buffer {
   void* ptr_;
 
  public:
-  explicit Buffer(void* ptr) : ptr_(ptr) {};
+  Buffer(void* ptr) : ptr_(ptr) {};
 
   // Get the raw data pointer from the buffer
   void* raw_ptr();
@@ -32,10 +32,14 @@ Buffer malloc(size_t size);
 
 void free(Buffer buffer);
 
+// Wait for running tasks to finish and free up memory
+// if allocation fails
+Buffer malloc_or_wait(size_t size);
+
 class Allocator {
   /** Abstract base class for a memory allocator. */
  public:
-  virtual Buffer malloc(size_t size) = 0;
+  virtual Buffer malloc(size_t size, bool allow_swap = false) = 0;
   virtual void free(Buffer buffer) = 0;
   virtual size_t size(Buffer buffer) const = 0;
 
@@ -49,4 +53,16 @@ class Allocator {
 
 Allocator& allocator();
 
+class CommonAllocator : public Allocator {
+  /** A general CPU allocator. */
+ public:
+  virtual Buffer malloc(size_t size, bool allow_swap = false) override;
+  virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
+
+ private:
+  CommonAllocator() = default;
+  friend Allocator& allocator();
+};
+
 } // namespace mlx::core::allocator