Specify exact package versions

.github/workflows/test.yaml

@@ -13,19 +13,41 @@ permissions:
   contents: read  # to fetch code
 
 jobs:
-  build:
+  built-latest:
     name: ${{ matrix.os }} Python ${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: ["ubuntu-latest"]
         python-version: ["3.9", "3.10", "3.11", "3.12"]
         include:
-        - os: macOS-11
-          python-version: "3.11"
         - os: windows-2019
           python-version: "3.11"
 
+    steps:
+    - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
+      with:
+        submodules: true
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # ratchet:actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -U .[latest,dev]
+    - name: Run tests
+      run: |
+        pytest -n auto jax_ml_stack
+
+  build:
+    name: ${{ matrix.os }} Python ${{ matrix.python-version }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version: ["3.10", "3.11", "3.12"]
+
     steps:
     - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
       with:

.gitignore

@@ -1,2 +1,3 @@
 *.egg-info
 Pipfile
+build/

pyproject.toml

@@ -15,9 +15,10 @@ keywords = []
 
 # pip dependencies of the project
 dependencies = [
-    "jax",
-    "flax",
-    "optax",
+    "jax==0.4.31",
+    "flax==0.8.5",
+    "optax==0.2.3",
+    "orbax==0.1.9",
 ]
 
 [project.urls]
@@ -33,6 +34,13 @@ dev = [
     "pytest-xdist",
 ]
 
+latest = [
+    "jax",
+    "flax",
+    "optax",
+    "orbax",
+]
+
 [tool.pyink]
 # Formatting configuration to follow Google style-guide
 line-length = 80