name: 'Build and Test with CUDA'
description: 'Build and test MLX with CUDA'

inputs:
  toolkit:
    description: 'The CUDA toolkit'
    required: true

runs:
  using: "composite"
  steps:
    - name: Install Python package
      shell: bash
      env:
        DEBUG: 1
        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
      run: pip install --no-build-isolation -e ".[dev]" -v

    - name: Build CPP only
      shell: bash
      run: |
        cmake . -B build \
          -DMLX_BUILD_CUDA=ON \
          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc \
          -DCMAKE_BUILD_TYPE=DEBUG
        cmake --build build -j $(nproc)