From 94235fc9cd2edd0f60e49bd079bdd71a5b5f8c10 Mon Sep 17 00:00:00 2001
From: "Lee J.O'Riordan" <lee@xanadu.ai>
Date: Wed, 15 Mar 2023 11:40:14 -0400
Subject: [PATCH] Add support for Github Actions build of jaxlib on Winx64

---
 .github/workflows/wheel_win_x64.yml | 85 +++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 .github/workflows/wheel_win_x64.yml

diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
new file mode 100644
index 000000000000..7fcfb73ab790
--- /dev/null
+++ b/.github/workflows/wheel_win_x64.yml
@@ -0,0 +1,85 @@
+name: Wheel::Windows::x86_64
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+  release:
+    types: [published]
+
+env:
+  DISTUTILS_USE_SDK: 1
+  MSSdk: 1
+
+jobs:
+  win-wheels:
+    strategy:
+      fail-fast: true
+      matrix:
+        os: [windows-2019]
+        arch: [AMD64]
+        pyver: ['3.8', '3.9', '3.10', '3.11']
+    name: ${{ matrix.os }} jaxlib Wheel-builder
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Cancel Previous Runs
+        uses: styfle/cancel-workflow-action@0.11.0
+        with:
+          access_token: ${{ github.token }}
+
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.pyver }}
+          cache: 'pip'
+
+      - name: Build wheels
+        env:
+          BAZEL_VC: "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC"
+        run: |
+          python -m pip install -r build/test-requirements.txt
+          "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
+
+          python.exe build\build.py `
+            --noenable_tpu `
+            --noenable_remote_tpu `
+            --noenable_mkl_dnn `
+            --target_cpu_features release `
+            --noenable_plugin_device `
+            --noremote_build `
+            --noenable_cuda `
+            --configure_only
+                    
+          bazel build --jobs=1 --subcommands `
+            --cxxopt="/Zm2000" `
+            --verbose_failures `
+            --jobs=1 `
+            --discard_analysis_cache `
+            --notrack_incremental_state `
+            --nokeep_state_after_build `
+            --experimental_local_memory_estimate="true" `
+            --local_ram_resources=2048 `
+            //jaxlib
+
+          bazel run `
+            --verbose_failures=true `
+            //build:build_wheel -- `
+            --output_path=$(Join-Path -Path (Get-Location) -ChildPath "\tmp") `
+            --cpu="AMD64"
+
+      - name: Run tests
+        env:
+          JAX_ENABLE_CHECKS: true
+          JAX_SKIP_SLOW_TESTS: true
+        run: |
+          python -m pip install -e ${{ github.workspace }}
+          python -m pip install --no-index --find-links ${{ github.workspace }}\dist jaxlib
+          echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
+          pytest -n auto --tb=short tests examples
+
+      - uses: actions/upload-artifact@v2
+        with:
+          name: jaxlib_wheels
+          path: ${{ github.workspace }}\dist\*.whl