Run cellfinder with JAX in Windows tests in CI

.github/workflows/test_and_deploy.yml

@@ -46,6 +46,8 @@ jobs:
         include:
         - os: macos-latest
           python-version: "3.10"
+        - os: windows-latest
+          python-version: "3.10"
 
     steps:
       # Cache the Keras model so we don't have to remake it every time
@@ -95,6 +97,8 @@ jobs:
     needs: [linting, manifest]
     name: Run brainmapper tests to check for breakages
     runs-on: ubuntu-latest
+    env:
+      KERAS_BACKEND: jax
     steps:
       - name: Cache Keras model
         uses: actions/cache@v3
@@ -115,10 +119,8 @@ jobs:
       - name: Install test dependencies
         run: |
           python -m pip install --upgrade pip wheel
-          # Install cellfinder from the latest SHA on this branch
-          python -m pip install git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA
-          # Install tensorflow as keras' default backend
-          python -m pip install "tf-nightly==2.16.0.dev20240101"
+          # Install cellfinder from the latest SHA on this branch (Keras with JAX backend)
+          python -m pip install "cellfinder[jax] @ git+$GITHUB_SERVER_URL/$GITHUB_REPOSITORY@$GITHUB_SHA"
           # Install checked out copy of brainglobe-workflows
           python -m pip install .[dev]
 

cellfinder/core/train/train_yml.py

@@ -338,7 +338,10 @@ def run(
 
     ensure_directory_exists(output_dir)
     model_weights = prep_model_weights(
-        install_path, model_weights, model, n_free_cpus
+        model_weights=model_weights,
+        install_path=install_path,
+        model_name=model,
+        n_free_cpus=n_free_cpus,
     )
 
     yaml_contents = parse_yaml(yaml_file)

pyproject.toml

@@ -154,4 +154,6 @@ passenv =
     XAUTHORITY
     NUMPY_EXPERIMENTAL_ARRAY_FUNCTION
     PYVISTA_OFF_SCREEN
+platform =
+    tf: linux|darwin  # skip TF backend on windows
 """