Add arm64 windows support (#678)

* add arm64 windows to python.yml

Also, added push event trigger to Python workflow to test on fork.

* add arm64 windows handling in tests

* Fix quotes

* fix typo

* Fix typo

* use different url for arm64

* Update condition for installing jax and flax

* run arm64 on limited set of tests

* Rename test step for Windows arm64

* add arm64 build

* Modify target condition for windows arm64

* add additional tests

Add specific test files to the testing workflow.

* update arm64 tests

* run all but tf tests

* change conditional, move install

* change test command sytax

to match previous syntax

* remove push trigger

Remove push trigger from Python workflow

Author: finnagin

.github/workflows/python-release.yml

@@ -94,6 +94,8 @@ jobs:
             target: x64
           - runner: windows-latest
             target: x86
+          - runner: windows-11-arm
+            target: arm64
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -103,7 +105,7 @@ jobs:
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
-          target: ${{ matrix.platform.target }}
+          target: ${{ matrix.platform.target == 'arm64' && 'aarch64-pc-windows-msvc' || matrix.platform.target }}
           args: --release --out dist --manifest-path bindings/python/Cargo.toml
           sccache: 'true'
       - name: Upload wheels

.github/workflows/python.yml

@@ -40,6 +40,12 @@ jobs:
               python: "3.12"
               numpy: numpy
               arch: "arm64"
+          - os: windows-11-arm
+            version:
+              torch: torch
+              python: "3.12"
+              numpy: numpy
+              arch: "arm64"
     defaults:
       run:
         working-directory: ./bindings/python
@@ -81,7 +87,7 @@ jobs:
       #     pip install -U pip
 
       - name: Install (torch)
-        if: matrix.version.arch != 'x64-freethreaded'
+        if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
         run: |
           pip install ${{ matrix.version.numpy }}
           pip install ${{ matrix.version.torch }}
@@ -93,22 +99,29 @@ jobs:
           pip install ${{ matrix.version.torch }} --index-url https://download.pytorch.org/whl/cu126
         shell: bash
 
+      - name: Install (torch windows arm64)
+        if: matrix.os == 'windows-11-arm'
+        run: |
+          pip install ${{ matrix.version.numpy }}
+          pip install ${{ matrix.version.torch }} --index-url https://download.pytorch.org/whl/cpu
+        shell: bash
+
       - name: Install (hdf5 non windows)
         if: matrix.os == 'ubuntu-latest' && matrix.version.arch != 'x64-freethreaded'
         run: |
           sudo apt-get update
           sudo apt-get install libhdf5-dev
 
       - name: Install (tensorflow)
-        if: matrix.version.arch != 'x64-freethreaded'
+        if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
         run: |
           pip install .[tensorflow]
           # Force reinstall of numpy, tensorflow uses numpy 2 even on 3.9
           pip install ${{ matrix.version.numpy }}
         shell: bash
 
       - name: Install (jax, flax)
-        if: matrix.os != 'windows-latest' && matrix.version.arch != 'x64-freethreaded'
+        if: runner.os != 'Windows' && matrix.version.arch != 'x64-freethreaded'
         run: 
           pip install .[jax]
         shell: bash
@@ -125,12 +138,19 @@ jobs:
           ruff format --check .
 
       - name: Run tests
-        if: matrix.version.arch != 'x64-freethreaded'
+        if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
         run: |
           cargo test
           pip install ".[testing]"
           pytest -sv tests/
 
+      - name: Run tests (Windows arm64)
+        if: matrix.os == 'windows-11-arm'
+        run: |
+          cargo test
+          pip install ".[testing]"
+          pytest -sv tests/ --ignore=tests/test_tf_comparison.py
+
       - name: Run tests (freethreaded)
         if: matrix.version.arch == 'x64-freethreaded'
         run: |