Build but not test CUDA arm64

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

@@ -64,9 +64,13 @@ runs:
         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 \
+        sudo apt-get install -y ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
-            libnccl2 libnccl-dev \
+
-            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
+    - name: Install NCCL
+      # For some reason NVIDIA did not provide arm64 package for libnccl2.
+      if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch != 'arm64' }}
+      shell: bash
+      run: sudo apt-get install -y libnccl2 libnccl-dev
 
     - name: CUDA packages and driver report
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}

.github/workflows/pull_request.yml

@@ -25,19 +25,32 @@ jobs:
   linux_build_and_test:
     needs: check_lint
     strategy:
-      matrix:
-        runner:
-          - ubuntu-22.04
-          - ubuntu-22.04-arm
       fail-fast: false
-    runs-on: ${{ matrix.runner }}
+      matrix:
+        arch: ['x86_64', 'aarch64']
+        toolkit: ['cpu', 'cuda-12.6', 'cuda-12.9']
+    # Use gpu runner for cuda/x86_64 builds, otherwise cpu runners.
+    runs-on: >-
+      ${{ (startsWith(matrix.toolkit, 'cuda') && matrix.arch == 'x86_64' && github.repository == 'ml-explore/mlx') &&
+          'gpu-t4-4-core' ||
+          (matrix.arch == 'aarch64' && 'ubuntu-22.04-arm' || 'ubuntu-22.04') }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
-      - uses: ./.github/actions/build-linux
-      - uses: ./.github/actions/test-linux
         with:
-          cpu-only: true
+          toolkit: ${{ matrix.toolkit }}
+      - uses: ./.github/actions/build-linux
+        with:
+          toolkit: ${{ matrix.toolkit }}
+      - uses: ./.github/actions/test-linux
+        # Skip tests for CUDA if:
+        # 1. building for aarch64;
+        # 2. not on offical repo.
+        if: >-
+          ${{ !(startsWith(matrix.toolkit, 'cuda') &&
+                (matrix.arch == 'aarch64' || github.repository != 'ml-explore/mlx')) }}
+        with:
+          cpu-only: ${{ matrix.toolkit == 'cpu' }}
 
   mac_build_and_test:
     if: github.repository == 'ml-explore/mlx'
@@ -53,24 +66,6 @@ jobs:
       - uses: ./.github/actions/setup-macos
       - uses: ./.github/actions/build-macos
 
-  cuda_build_and_test:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      fail-fast: false
-      matrix:
-        toolkit: ['cuda-12.6', 'cuda-12.9']
-    runs-on: gpu-t4-4-core
-    needs: check_lint
-    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
-
   build_documentation:
     if: github.repository == 'ml-explore/mlx'
     runs-on: ubuntu-22.04