name: Setup

inputs:
  python-version:
    required: false
    default: '3.9'
  torch-version:
    required: false
    default: '2.0.0'
  cuda-version:
    required: false
    default: cpu
  full_install:
    required: false
    default: true

runs:
  using: composite

  steps:
    - name: Set up Python ${{ inputs.python-version }}
      uses: actions/setup-python@v4.3.0
      with:
        python-version: ${{ inputs.python-version }}
        check-latest: true
        cache: pip
        cache-dependency-path: |
          pyproject.toml

    - name: Install PyTorch ${{ inputs.torch-version }}+${{ inputs.cuda-version }}
      run: |
        pip install torch==${{ inputs.torch-version }} --extra-index-url https://download.pytorch.org/whl/${{ inputs.cuda-version }}
        python -c "import torch; print('PyTorch:', torch.__version__)"
        python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
        python -c "import torch; print('CUDA:', torch.version.cuda)"
      shell: bash

    - name: Install extension packages
      if: ${{ inputs.full_install == 'true' }}
      run: |
        pip install scipy
        pip install --no-index --upgrade torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${{ inputs.torch-version }}+${{ inputs.cuda-version }}.html
      shell: bash