From ef0a73024acc2b5e257207dffee419418f7160e1 Mon Sep 17 00:00:00 2001
From: James Spencer <jamessspencer@google.com>
Date: Mon, 9 Dec 2024 09:59:24 +0000
Subject: [PATCH] Simplify and update github CI

---
 .github/workflows/ci-build.yaml | 38 ++++++---------------------------
 setup.py                        |  4 ++--
 2 files changed, 8 insertions(+), 34 deletions(-)

diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
index a589840..e58795d 100644
--- a/.github/workflows/ci-build.yaml
+++ b/.github/workflows/ci-build.yaml
@@ -16,26 +16,14 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ['3.11']
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Get pip cache dir
-      id: pip-cache
-      run: |
-        python -m pip install --upgrade pip
-        echo "::set-output name=dir::$(pip cache dir)"
-    - name: pip cache
-      uses: actions/cache@v2
-      with:
-        path: ${{ steps.pip-cache.outputs.dir }}
-        key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
         pip install -e '.[testing]'
@@ -63,23 +51,9 @@ jobs:
             os: ubuntu-latest
             package-overrides: "none"
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Get pip cache dir
-      id: pip-cache
-      run: |
-        python -m pip install --upgrade pip
-        echo "::set-output name=dir::$(pip cache dir)"
-    - name: pip cache
-      uses: actions/cache@v2
-      with:
-        path: ${{ steps.pip-cache.outputs.dir }}
-        key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
+      uses: actions/setup-python@v5
     - name: Install dependencies
       run: |
         pip install -e '.[testing]'
@@ -90,6 +64,6 @@ jobs:
       run: |
         python -m pytest
     - name: Run multi-device tests
-      if: matrix.python-version == '3.10'
+      if: matrix.python-version == '3.11'
       run: |
         FERMINET_CHEX_N_CPU_DEVICES=2 python -m pytest ferminet/tests/train_test.py
diff --git a/setup.py b/setup.py
index 80e785a..aa85799 100644
--- a/setup.py
+++ b/setup.py
@@ -24,11 +24,11 @@
     'attrs',
     'chex',
     'h5py',
-    'folx @ git+https://github.com/microsoft/folx',
+    'folx @ git+https://github.com/microsoft/folx@main',
     'jax',
     'jaxlib',
     # TODO(b/230487443) - use released version of kfac.
-    'kfac_jax @ git+https://github.com/deepmind/kfac-jax',
+    'kfac_jax @ git+https://github.com/google-deepmind/kfac-jax@main',
     'ml-collections',
     'optax',
     'numpy',