name: 'Build Linux wheel'
description: 'Build Linux wheel'

inputs:
  build-backend:
    description: 'Build the backend mlx-cpu package'
    type: boolean
    required: false
    default: false
  arch:
    description: 'Platform architecture tag'
    required: true
    type: choice
    options:
      - x86_64
      - aarch64

runs:
  using: "composite"
  steps:
    - name: Build MLX
      shell: bash
      run: pip install -e . -v

    - name: Build Python package
      shell: bash
      run: |
        pip install auditwheel patchelf build
        python setup.py clean --all
        MLX_BUILD_STAGE=1 python -m build -w
        auditwheel repair dist/mlx-*.whl \
          --plat manylinux_2_35_${{ inputs.arch }} \
          --exclude libmlx.so* \
          --only-plat

    - name: Build backend package
      if: ${{ inputs.build-backend }}
      shell: bash
      run: |
        python setup.py clean --all
        MLX_BUILD_STAGE=2 python -m build -w
        auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_${{ inputs.arch }}
