Merge pull request #22 from rdyro/add-matrix-to-ci

adding matrix to the ci
rdyro · Dec 2, 2024 · 7987b1d · 7987b1d
2 parents f198118 + c3ffffd
commit 7987b1d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml
@@ -20,19 +20,26 @@ concurrency:
 
 jobs:
   lint_and_typecheck:
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
     runs-on: ubuntu-latest
     timeout-minutes: 5
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python 3.11
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11
+          python-version: ${{ matrix.python-version }}
       - run: |
           python3 -m pip install -U ruff
           ruff format --diff torch2jax
           ruff check torch2jax
   pytest:
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+        jax-version: ["0.4.33", "0.4.35", "latest"]
     runs-on: ubuntu-latest
     timeout-minutes: 5
     steps:
@@ -41,10 +48,10 @@ jobs:
         run: |
           sudo DEBIAN_FRONTEND=noninteractive apt update -y
           sudo DEBIAN_FRONTEND=noninteractive apt install -y ninja-build coreutils
-      - name: Set up Python 3.11
+      - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11
+          python-version: ${{ matrix.python-version }}
       - name: Get uv cache dir
         id: uv-cache
         run: |
@@ -61,8 +68,11 @@ jobs:
           [[ ! -d ~/.venv ]] && uv venv --seed ~/.venv
           source ~/.venv/bin/activate
           uv pip install pytest && pip install pytest
-          uv pip install -U jax torch
-          uv pip install -U .
+          if [[ ${{ matrix.jax-version }} == "latest" ]]; then
+            uv pip install -U torch jax .
+          else
+            uv pip install torch jax==${{ matrix.jax-version }} .
+          fi
       - name: Get extension path
         id: cpp-extension-cache
         run: |

diff --git a/scripts/make_aliased_package.py b/scripts/make_aliased_package.py
@@ -77,12 +77,12 @@
     # rewrite the pyproject.toml file
     pyproject_toml = Path(root_path / "pyproject.toml").read_text()
     new_pyproject_toml = re.sub(
-        f"name\s*=\s*\"{current_config['project']['name']}\"",
+        f"name\\s*=\\s*\"{current_config['project']['name']}\"",
         f'name = "{new_name}"',
         pyproject_toml,
     )
     new_pyproject_toml = re.sub(
-        f"{current_config['project']['name']}\s*=", f"{new_name} = ", new_pyproject_toml
+        f"{current_config['project']['name']}\\s*=", f"{new_name} = ", new_pyproject_toml
     )
     Path(new_package_path / "pyproject.toml").write_text(new_pyproject_toml)