Split out indexed tests into a separate group

.github/workflows/lint.yml

@@ -19,9 +19,11 @@ jobs:
         with:
           enable-cache: true
 
-      - name: Install pandoc
-        run: |
-          sudo apt install -y pandoc
+      - name: Cache virtual environment
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ hashFiles('uv.lock') }}
 
       - name: Install dependencies
         run: |

.github/workflows/test.yml

@@ -14,23 +14,43 @@ jobs:
     strategy:
       matrix:
         python-version: ["'<3.13'", "'3.13'"]
+        test-group:
+          - name: indexed
+            paths: tests/test_handlers_indexed*.py
+          - name: torch
+            paths: tests/test_handlers_torch*.py
+          - name: pyro
+            paths: tests/test_handlers_pyro*.py
+          - name: jax
+            paths: tests/test_handlers_jax*.py
+          - name: numpyro
+            paths: tests/test_handlers_numpyro*.py
+          - name: llm
+            paths: tests/test_handlers_llm*.py
+          - name: examples
+            paths: tests/test_examples_*.py
+          - name: core
+            paths: effectful/ tests/test_ops_*.py tests/test_internals_*.py
+    name: "${{ matrix.test-group.name }} (py ${{ matrix.python-version }})"
     steps:
       - uses: actions/checkout@v4
-      
+
       - name: Install uv
         uses: astral-sh/setup-uv@v5
         with:
           enable-cache: true
 
-      - name: Install pandoc
-        run: |
-          sudo apt install -y pandoc
+      - name: Cache virtual environment
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
 
       - name: Install Python dependencies
         run: |
           uv sync --all-extras --dev --python ${{ matrix.python-version }}
-          
+
       - name: Run tests
         env:
           MKL_NUM_THREADS: 1
-        run: uv run ./scripts/test.sh
+        run: uv run pytest ${{ matrix.test-group.paths }} -n auto

.github/workflows/test_docs.yml

@@ -23,6 +23,12 @@ jobs:
         run: |
           sudo apt install -y pandoc
 
+      - name: Cache virtual environment
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ hashFiles('uv.lock') }}
+
       - name: Install Python dependencies
         run: |
           uv sync --all-extras --dev

.github/workflows/test_llm.yml

@@ -22,6 +22,12 @@ jobs:
         with:
           enable-cache: true
 
+      - name: Cache virtual environment
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
+
       - name: Install Python dependencies
         run: |
           uv sync --all-extras --dev --python ${{ matrix.python-version }}

.github/workflows/test_notebooks.yml

@@ -30,6 +30,12 @@ jobs:
         run: |
           sudo apt install -y pandoc
 
+      - name: Cache virtual environment
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ hashFiles('uv.lock') }}
+
       - name: Install Python packages
         run: |
           uv sync --all-extras --dev