Linux on arm (#2751)

* try linux on arm

* ssh

* fix

.circleci/config.yml

@@ -1,579 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
-
-parameters:
-  nightly_build:
-    type: boolean
-    default: false
-  test_release:
-    type: boolean
-    default: false
-
-jobs:
-  build_documentation:
-    parameters:
-      upload-docs:
-        type: boolean
-        default: false
-    macos:
-      xcode: "26.0.0"
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            brew install python@3.10
-            brew install doxygen
-            python3.10 -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install -r docs/requirements.txt
-            pip install . -v
-      - when:
-          condition:
-            not: << parameters.upload-docs >>
-          steps:
-            - run:
-               name: Build documentation
-               command: |
-                 source env/bin/activate
-                 cd docs && doxygen && make html O=-W
-      - when:
-          condition: << parameters.upload-docs >>
-          steps:
-            - add_ssh_keys:
-                fingerprints:
-                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
-            - run:
-               name: Upload documentation
-               command: |
-                 source env/bin/activate
-                 git config user.email "mlx@group.apple.com"
-                 git config user.name "CircleCI Docs"
-                 git checkout gh-pages
-                 git rebase main
-                 cd docs
-                 git rm -rf build/html
-                 doxygen && make html O=-W
-                 git add -f build/html
-                 git commit -m "rebase"
-                 git push -f origin gh-pages
-
-  linux_build_and_test:
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Run style checks
-          command: |
-            pip install pre-commit
-            pre-commit run --all
-            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
-      - run:
-          name: Install dependencies
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            python -m unittest discover python/tests -v
-            mpirun --bind-to none -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-            if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build
-            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-            make -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests
-
-  mac_build_and_test:
-    parameters:
-      xcode_version:
-        type: string
-        default: "26.0.0"
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    resource_class: m4pro.medium
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
-              brew install openmpi uv
-      - run:
-          name: Install Python package
-          command: |
-            uv venv --python 3.10
-            uv pip install \
-              nanobind==2.4.0 \
-              cmake \
-              numpy \
-              torch \
-              tensorflow \
-              unittest-xml-reporting
-            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              uv pip install -e . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            uv pip install typing_extensions
-            uv run --no-project setup.py generate_stubs
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
-            mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
-            mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
-            if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-      - run:
-          name: Build example extension
-          command: |
-            source .venv/bin/activate
-            cd examples/extensions
-            uv pip install -r requirements.txt
-            uv run --no-project setup.py build_ext --inplace
-            uv run --no-project python test.py
-      - store_test_results:
-          path: test-results
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
-      - run:
-          name: Run CPP tests
-          command: |
-            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
-      - run:
-          name: Build small binary
-          command: |
-            source .venv/bin/activate
-            cd build/
-            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
-              -DBUILD_SHARED_LIBS=ON \
-              -DMLX_BUILD_CPU=OFF \
-              -DMLX_BUILD_SAFETENSORS=OFF \
-              -DMLX_BUILD_GGUF=OFF \
-              -DMLX_METAL_JIT=ON
-            make -j `sysctl -n hw.ncpu`
-      - run:
-          name: Run Python tests with JIT
-          command: |
-            CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-              uv pip install -e . -v
-            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
-              METAL_DEBUG_ERROR_MODE=0 \
-              uv run --no-project python -m xmlrunner discover \
-                -v python/tests \
-                -o test-results/gpu_jit
-
-  cuda_build_and_test:
-    parameters:
-      image_date:
-        type: string
-        default: "2023.11.1"
-    machine:
-      image: "linux-cuda-12:<< parameters.image_date >>"
-      resource_class: gpu.nvidia.small.gen2
-    steps:
-      - checkout
-      - restore_cache:
-          keys:
-            - cuda-<< parameters.image_date >>-{{ arch }}-
-      - run:
-          name: Install dependencies
-          command: |
-            sudo apt-get update
-            sudo apt-get install libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install libnccl2 libnccl-dev
-            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
-            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
-            rm -rf ccache-4.11.3-linux-x86_64
-            curl -LsSf https://astral.sh/uv/install.sh | sh
-      - run:
-          name: Set CCache size
-          command: ccache --max-size 1G
-      - run:
-          name: Install Python package
-          command: |
-            uv venv
-            uv pip install cmake
-            DEBUG=1 CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              uv pip install -e ".[dev]" -v
-      - run:
-          name: Run Python tests
-          command: |
-            source .venv/bin/activate
-            LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
-            LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
-      - run:
-          name: Build CPP only
-          command: |
-            source .venv/bin/activate
-            cmake . -B build \
-              -DMLX_BUILD_CUDA=ON \
-              -DCMAKE_CUDA_COMPILER=`which nvcc` \
-              -DCMAKE_BUILD_TYPE=DEBUG
-            cmake --build build -j `nproc`
-      - run:
-          name: Run CPP tests
-          command: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-      - run:
-          name: CCache report
-          command: |
-            ccache --show-stats
-            ccache --zero-stats
-            ccache --cleanup
-      - save_cache:
-          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
-          paths:
-            - /home/circleci/.cache/ccache
-
-  build_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.10"
-      xcode_version:
-        type: string
-        default: "26.0.0"
-      build_env:
-        type: string
-        default: ""
-      macosx_deployment_target:
-        type: string
-        default: ""
-    macos:
-      xcode: << parameters.xcode_version >>
-    resource_class: m4pro.medium
-    environment:
-      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            xcodebuild -downloadComponent MetalToolchain
-            mkdir -p ~/miniconda3
-            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o ~/miniconda3/miniconda.sh
-            bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
-            rm ~/miniconda3/miniconda.sh
-            source ~/miniconda3/bin/activate
-            conda init --all
-            conda create -n env python=<< parameters.python_version >> -y
-            conda activate env
-            pip install --upgrade cmake
-            pip install nanobind==2.4.0
-            pip install --upgrade setuptools
-            pip install numpy
-            pip install twine
-            pip install build
-      - run:
-          name: Install Python package
-          command: |
-            conda activate env
-            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
-              pip install . -v
-      - run:
-          name: Generate package stubs
-          command: |
-            conda activate env
-            pip install typing_extensions
-            python setup.py generate_stubs
-      - run:
-          name: Build Python package
-          command: |
-            conda activate env
-            python setup.py clean --all
-            << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
-      - when:
-          condition:
-            equal: ["3.10", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  conda activate env
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 python -m build -w
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  conda activate env
-                  twine upload dist/*
-      - store_artifacts:
-          path: dist/
-
-  build_linux_release:
-    parameters:
-      python_version:
-        type: string
-        default: "3.10"
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: large
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            PYTHON=python<< parameters.python_version >>
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            sudo apt-get update
-            TZ=Etc/UTC sudo apt-get -y install tzdata
-            sudo add-apt-repository -y ppa:deadsnakes/ppa
-            sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
-            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            $PYTHON -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            << parameters.build_env >> pip install ".[dev]" -v
-            pip install typing_extensions
-            python setup.py generate_stubs
-            python setup.py clean --all
-            MLX_BUILD_STAGE=1 << parameters.build_env >> python -m build -w
-            bash python/scripts/repair_linux.sh
-      - when:
-          condition:
-            equal: ["3.10", << parameters.python_version >>]
-          steps:
-            - run:
-                name: Build common package
-                command: |
-                  source env/bin/activate
-                  python setup.py clean --all
-                  << parameters.build_env >> MLX_BUILD_STAGE=2 \
-                    python -m build -w
-                  auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload packages
-                command: |
-                  source env/bin/activate
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-  build_cuda_release:
-    parameters:
-      build_env:
-        type: string
-        default: ""
-    machine:
-      image: ubuntu-2204:current
-      resource_class: xlarge
-    steps:
-      - checkout
-      - run:
-          name: Build wheel
-          command: |
-            export DEBIAN_FRONTEND=noninteractive
-            export NEEDRESTART_MODE=a
-            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
-            sudo dpkg -i cuda-keyring_1.1-1_all.deb
-            sudo apt-get update
-            sudo apt-get install cuda-toolkit-12-9 libcudnn9-dev-cuda-12
-            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install zip
-            pip install auditwheel
-            pip install patchelf
-            pip install build
-            pip install twine
-            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
-            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-            << parameters.build_env >> MLX_BUILD_STAGE=2 \
-              CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              python -m build -w
-            bash python/scripts/repair_cuda.sh
-      - when:
-          condition: << parameters.build_env >>
-          steps:
-            - run:
-                name: Upload package
-                command: |
-                  twine upload wheelhouse/*.whl
-      - store_artifacts:
-          path: wheelhouse/
-
-workflows:
-  build_and_test:
-    when:
-      and:
-        - matches:
-            pattern: "^(?!pull/)[-\\w]+$"
-            value: << pipeline.git.branch >>
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - mac_build_and_test:
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test
-      - cuda_build_and_test:
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-      - build_documentation 
-
-  build_pypi_release:
-    when:
-      and:
-        - not: << pipeline.parameters.nightly_build >>
-        - not: << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["PYPI_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_documentation:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          upload-docs: true
-      - build_linux_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              build_env: ["PYPI_RELEASE=1"]
-      - build_cuda_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-          matrix:
-            parameters:
-              build_env: ["PYPI_RELEASE=1"]
-
-  prb:
-    when:
-      matches:
-        pattern: "^pull/\\d+(/head)?$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - hold:
-          type: approval
-      - apple/authenticate:
-          context: pr-approval
-      - mac_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              macosx_deployment_target: ["13.5", "15.0"]
-      - linux_build_and_test:
-          requires: [ hold ]
-      - cuda_build_and_test:
-          requires: [ hold ]
-          matrix:
-            parameters:
-              image_date: ["2023.11.1", "2025.05.1"]
-  nightly_build:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.nightly_build >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-      - build_cuda_release
-
-  build_dev_release:
-    when:
-      and:
-        - equal: [ main, << pipeline.git.branch >> ]
-        - << pipeline.parameters.test_release >>
-    jobs:
-      - build_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              macosx_deployment_target: ["13.5", "14.0", "15.0"]
-              build_env: ["DEV_RELEASE=1"]
-              xcode_version: ["26.0.0"]
-      - build_linux_release:
-          matrix:
-            parameters:
-              python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-              build_env: ["DEV_RELEASE=1"]
-      - build_cuda_release:
-          matrix:
-            parameters:
-              build_env: ["DEV_RELEASE=1"]

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

@@ -12,13 +12,9 @@ runs:
     - name: Build package
       shell: bash
       env:
-        MLX_BUILD_STAGE: 2
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
       run: |
         pip install auditwheel build patchelf setuptools
         python setup.py clean --all
-        python -m build -w
-
-        if [ -f "python/scripts/repair_cuda.sh" ]; then
-          bash python/scripts/repair_cuda.sh
-        fi
+        MLX_BUILD_STAGE=2 python -m build -w
+        bash python/scripts/repair_cuda.sh

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

@@ -2,19 +2,10 @@ name: 'Build and Test with CUDA'
 description: 'Build and test MLX with CUDA'
 
 inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
   nvcc-location:
     description: 'Location of nvcc compiler'
     required: true
     default: '/usr/local/cuda-12.9/bin/nvcc'
-    # this value is dependent on the CUDA tools installed in the setup-linux workflow
 
 runs:
   using: "composite"
@@ -27,7 +18,6 @@ runs:
       run: pip install -e ".[dev]" -v
 
     - name: Run Python tests - CPU
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -35,7 +25,6 @@ runs:
       run: python -m unittest discover python/tests -v
 
     - name: Run Python tests - GPU
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -43,7 +32,6 @@ runs:
       run: python -m tests discover python/tests -v
 
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         cmake . -B build \
@@ -53,12 +41,5 @@ runs:
         cmake --build build -j $(nproc)
     
     - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: bash
       run: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
-
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-cuda-release
-      with:
-        nvcc-location: ${{ inputs.nvcc-location }}

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

@@ -0,0 +1,40 @@
+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,33 +1,14 @@
 name: 'Build and Test on Linux'
 description: 'Build and test MLX on Linux'
 
-inputs:
-  build-type:
-    description: 'Build type'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-    type: boolean
-
 runs:
   using: "composite"
   steps:
-    - name: Set DEBUG
-      shell: sh
-      if: inputs.build-type == 'debug'
-      run: echo "DEBUG=1" >> $GITHUB_ENV
-
     - name: Install Python package
       shell: sh
       env:
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
+        DEBUG: 1
       run: pip install -e ".[dev]" -v
     
     - name: Generate package stubs
@@ -37,7 +18,6 @@ runs:
         python setup.py generate_stubs
     
     - name: Run Python tests
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         python -m unittest discover python/tests -v
@@ -50,7 +30,6 @@ runs:
         fi
     
     - name: Build CPP only
-      if: inputs.build-type == 'debug'
       shell: bash
       run: |
         mkdir -p build && cd build
@@ -58,21 +37,5 @@ runs:
         make -j $(nproc)
     
     - name: Run CPP tests
-      if: ${{ inputs.build-type == 'debug' && inputs.run-tests == 'true' }}
       shell: sh
       run: ./build/tests/tests
-    
-    - name: Build Python package
-      if: inputs.build-type == 'release'
-      shell: bash
-      run: |
-        pip install auditwheel patchelf build
-        python setup.py clean --all
-        MLX_BUILD_STAGE=1 python -m build -w
-        if [ -f "python/scripts/repair_linux.sh" ]; then
-          bash python/scripts/repair_linux.sh
-        fi
-
-        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
-        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64

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

@@ -6,11 +6,16 @@ inputs:
     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(s)
+    - name: Build Python package
       shell: bash
       env:
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
@@ -18,5 +23,11 @@ runs:
         uv pip install build
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=1 uv run -m build -w
+    - name: Build backend package
+      if: ${{ inputs.build-backend }}
+      shell: bash
+      env:
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
+      run: |
         uv run --no-project setup.py clean --all
         MLX_BUILD_STAGE=2 uv run -m build -w

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

@@ -1,20 +1,6 @@
 name: 'Build and Test on macOS'
 description: 'Build and test MLX on macOS'
 
-inputs:
-  build-type:
-    description: 'Build type (debug, release)'
-    required: false
-    default: 'debug'
-    type: choice
-    options:
-      - debug
-      - release
-  run-tests:
-    description: 'Whether to run tests'
-    required: false
-    default: 'true'
-
 runs:
   using: "composite"
   steps:
@@ -22,7 +8,7 @@ runs:
       shell: sh
       env:
         DEBUG: 1
-        DEV_RELEASE: 1
+        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
       run: |
         uv pip install --upgrade pip
         uv pip install cmake setuptools nanobind==2.4.0
@@ -35,13 +21,11 @@ runs:
         uv run --no-project setup.py generate_stubs
 
     - name: Install tests dependencies
-      if: inputs.run-tests == 'true'
       shell: sh
       run: |
         uv pip install numpy torch tensorflow unittest-xml-reporting
 
     - name: Run Python tests
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -53,7 +37,6 @@ runs:
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
     
     - name: Build example extension
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         cd examples/extensions
@@ -62,7 +45,6 @@ runs:
         uv run --no-project test.py
     
     - name: Build CPP only
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         mkdir -p build
@@ -71,7 +53,6 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run CPP tests
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         DEVICE: gpu
@@ -80,7 +61,6 @@ runs:
       run: ./build/tests/tests
     
     - name: Build small binary with JIT
-      if: inputs.run-tests == 'true'
       shell: bash
       run: |
         mkdir -p build
@@ -94,7 +74,6 @@ runs:
         make -j $(sysctl -n hw.ncpu)
     
     - name: Run Python tests with JIT
-      if: inputs.run-tests == 'true'
       shell: bash
       env:
         LOW_MEMORY: 1
@@ -107,14 +86,3 @@ runs:
         uv run -m xmlrunner discover \
             -v python/tests \
             -o test-results/gpu_jit
-
-    - name: Build macOS 14 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 14.0
-    - name: Build macOS 15 package
-      if: inputs.build-type == 'release'
-      uses: ./.github/actions/build-macos-release
-      with:
-        macos-target: 15.0

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

@@ -2,11 +2,6 @@ name: 'Setup macOS Environment'
 description: 'Install dependencies for macOS builds'
 
 inputs:
-  install-mpi:
-    description: 'Whether to install MPI'
-    required: false
-    default: 'true'
-    type: boolean
   python-version:
     description: 'Python version to use'
     required: false
@@ -17,7 +12,6 @@ runs:
   steps:
     - name: Install Homebrew packages
       shell: sh
-      if: inputs.install-mpi == 'true'
       run: /opt/homebrew/bin/brew install openmpi
     
     - name: Verify MetalToolchain installed

.github/workflows/nightly.yml

@@ -18,10 +18,10 @@ jobs:
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
-          run-tests: false
+          build-backend: ${{ matrix.python-version == '3.10' }}
+          arch: "x86_64"
       - name: Upload mlx artifacts
         uses: actions/upload-artifact@v5
         with:
@@ -41,7 +41,10 @@ jobs:
       fail-fast: false
       matrix:
         python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
+        runner:
+          - ubuntu-22.04
+          - ubuntu-22.04-arm
+    runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux
@@ -54,8 +57,6 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10", "3.13"]
-    env:
-      MACOSX_DEPLOYMENT_TARGET: "15.0"
     runs-on: [self-hosted, macos]
     steps:
       - uses: actions/checkout@v5
@@ -64,6 +65,17 @@ jobs:
           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_with_tests:
     if: github.repository == 'ml-explore/mlx'
     runs-on: gpu-t4-4-core

.github/workflows/pull_request.yml

@@ -14,7 +14,13 @@ jobs:
       - uses: pre-commit/action@v3.0.1
 
   linux_build_and_test:
-    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        runner:
+          - ubuntu-22.04
+          - ubuntu-22.04-arm
+      fail-fast: false
+    runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux

.github/workflows/release.yml

@@ -45,7 +45,12 @@ jobs:
     strategy:
       matrix:
         python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
-    runs-on: ubuntu-22.04
+        include:
+          - runner: ubuntu-24.04
+            arch: x64
+          - runner: ubuntu-24.04-arm64
+            arch: arm64
+    runs-on: ${{ matrix.runner }}
     env:
       PYPI_RELEASE: 1
     steps:
@@ -53,10 +58,10 @@ jobs:
       - uses: ./.github/actions/setup-linux
         with:
           python-version: ${{ matrix.python_version }}
-      - uses: ./.github/actions/build-linux
+      - uses: ./.github/actions/build-linux-release
         with:
-          build-type: release
-          run-tests: false
+          build-backend: ${{ matrix.python-version == '3.10' }}
+          arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v5
         with:
@@ -82,10 +87,27 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-      - uses: ./.github/actions/build-macos
+      - name: Install dependencies
+        shell: sh
+        run: |
+          uv pip install --upgrade pip
+          uv pip install cmake setuptools nanobind==2.4.0
+          uv pip install -e . -v
+      - name: Generate package stubs
+        shell: bash
+        run: |
+          uv pip install typing_extensions
+          uv run --no-project setup.py generate_stubs
+      - name: Build macOS 14 package
+        uses: ./.github/actions/build-macos-release
         with:
-          build-type: release
-          run-tests: false
+          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:

benchmarks/cpp/irregular_strides.cpp

@@ -75,7 +75,7 @@ void time_irregular_binary_ops_3D() {
 
 void time_irregular_binary_ops_4D() {
   auto device = mx::default_device();
-  std::vector<int> shape = {8, 8, 512, 512};
+  mx::Shape shape = {8, 8, 512, 512};
   auto a = mx::random::uniform(shape);
   auto b = mx::random::uniform(shape);
 
@@ -115,7 +115,7 @@ void time_irregular_binary_ops_4D() {
 
 void time_irregular_reshape() {
   auto device = mx::default_device();
-  std::vector<int> shape;
+  mx::Shape shape;
   auto reshape_fn = [&shape, device](const mx::array& a) {
     return mx::reshape(a, shape, device);
   };
@@ -170,7 +170,7 @@ void time_irregular_astype_1D() {
 void time_irregular_astype_2D() {
   auto device = mx::default_device();
   int size = 2048;
-  std::vector<int> shape = {size, size};
+  mx::Shape shape = {size, size};
 
   auto a = mx::random::uniform(shape);
   TIMEM("2D regular", mx::astype, a, mx::int32, device);

mlx/version.h

@@ -4,7 +4,7 @@
 
 #define MLX_VERSION_MAJOR 0
 #define MLX_VERSION_MINOR 29
-#define MLX_VERSION_PATCH 4
+#define MLX_VERSION_PATCH 5
 #define MLX_VERSION_NUMERIC \
   (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
 

python/scripts/repair_linux.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 auditwheel repair dist/* \
-  --plat manylinux_2_35_x86_64 \
+  --plat manylinux_2_35_${1} \
   --only-plat \
   --exclude libmlx* \
   -w wheel_tmp

python/src/ops.cpp

@@ -4335,7 +4335,7 @@ void init_ops(nb::module_& m) {
       nb::kw_only(),
       "stream"_a = nb::none(),
       nb::sig(
-          "def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', dtype: Optional[Dtype], *, stream: Union[None, Stream, Device] = None) -> array"),
+          "def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
         Dequantize the matrix ``w`` using quantization parameters.
 

tests/linalg_tests.cpp

@@ -425,9 +425,11 @@ TEST_CASE("test matrix pseudo-inverse") {
     const auto A = array({1.0, 2.0, 3.0, 4.0}, {2, 2});
     const auto A_pinv = linalg::pinv(A, Device::cpu);
     const auto A_again = matmul(matmul(A, A_pinv), A);
-    CHECK(allclose(A_again, A).item<bool>());
+    CHECK(allclose(A_again, A, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
     const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv);
-    CHECK(allclose(A_pinv_again, A_pinv).item<bool>());
+    CHECK(allclose(A_pinv_again, A_pinv, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
   }
   { // Rectangular matrix m < n
     const auto prng_key = random::key(42);
@@ -437,9 +439,11 @@ TEST_CASE("test matrix pseudo-inverse") {
     CHECK_FALSE(allclose(zeros, A_pinv, /* rtol = */ 0, /* atol = */ 1e-6)
                     .item<bool>());
     const auto A_again = matmul(matmul(A, A_pinv), A);
-    CHECK(allclose(A_again, A).item<bool>());
+    CHECK(allclose(A_again, A, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
     const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv);
-    CHECK(allclose(A_pinv_again, A_pinv).item<bool>());
+    CHECK(allclose(A_pinv_again, A_pinv, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
   }
   { // Rectangular matrix m > n
     const auto prng_key = random::key(10);
@@ -449,9 +453,11 @@ TEST_CASE("test matrix pseudo-inverse") {
     CHECK_FALSE(allclose(zeros2, A_pinv, /* rtol = */ 0, /* atol = */ 1e-6)
                     .item<bool>());
     const auto A_again = matmul(matmul(A, A_pinv), A);
-    CHECK(allclose(A_again, A).item<bool>());
+    CHECK(allclose(A_again, A, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
     const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv);
-    CHECK(allclose(A_pinv_again, A_pinv).item<bool>());
+    CHECK(allclose(A_pinv_again, A_pinv, /* rtol = */ 1e-5, /* atol = */ 1e-5)
+              .item<bool>());
   }
 }