rocm_jax/.github/workflows/windows_ci.yml

name: CI - Windows CPU
on:
  schedule:
    - cron: "0 12 * * *" # Daily at 12:00 UTC
  workflow_dispatch: # allows triggering the workflow run manually
  pull_request:
    types: [ labeled ] # allow force-windows-run label

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: true

env:
  DISTUTILS_USE_SDK: 1
  MSSdk: 1

jobs:
  win-wheels:
    if: ${{ (github.event.action != 'labeled') || (github.event.label.name == 'windows:force-run')}}
    strategy:
      fail-fast: true
      matrix:
        os: [windows-2019-32core]
        arch: [AMD64]
        pyver: ['3.10']
    name: Windows CI build
    runs-on: ${{ matrix.os }}

    steps:

      - name: Install LLVM/Clang
        run: choco install llvm --version=18.1.4 --yes --no-progress --allow-downgrade

      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
        with:
          path: jax

      - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38  # v5.4.0
        with:
          python-version: ${{ matrix.pyver }}
          cache: 'pip'

      - name: Build wheels
        env:
          BAZEL_VC: "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC"
          JAXLIB_NIGHTLY: true  # Tag the wheels as dev versions
        run: |
          cd jax
          python -m pip install uv~=0.5.30
          python -m uv pip install -r build/test-requirements.txt --upgrade numpy==2.0.0 scipy==1.13.1
          "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
          python.exe build\build.py build --wheels=jaxlib `
            --bazel_options=--color=yes `
            --bazel_options=--config=win_clang `
            --verbose

      - uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b  # v4.5.0
        with:
          name: wheels
          path: ${{ github.workspace }}\jax\dist\*.whl
          retention-days: 5

      - name: Run tests
        env:
          JAX_ENABLE_CHECKS: true
          JAX_SKIP_SLOW_TESTS: true
          PY_COLORS: 1
        run: |
          cd jax
          python -m uv pip install --pre --find-links ${{ github.workspace }}\jax\dist jaxlib \
            -e ${{ github.workspace }}\jax
          echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
          pytest -n auto --tb=short tests examples