Remove support for JAX 0.4 and 0.5. Fix for 0.6.

- Fix installation of jax in CI
- Display JAX version in CI

Author: pchanial

.github/workflows/ci.yml

@@ -29,15 +29,15 @@ jobs:
       matrix:
         python:
         - version: '3.10'
-          jax: '>=0.4,<0.5'
+          jax: '>=0.6,<0.7'
         - version: '3.11'
-          jax: '>=0.4,<0.5'
+          jax: '>=0.7,<0.8'
         - version: '3.12'
-          jax: '>=0.5,<0.6'
+          jax: '>=0.7,<0.8'
         - version: '3.13'
-          jax: '>=0.5,<0.6'
+          jax: '>=0.8,<0.9'
         - version: '3.14'
-          jax: '>=0.8'
+          jax: '>=0.9'
 
     steps:
       - name: Checkout code
@@ -63,7 +63,10 @@ jobs:
           uv pip install "jax${{ matrix.python.jax }}"
 
       - name: Run unit tests
-        run: uv run pytest -v --cov=zeropybench --cov-report=xml tests
+        run: |
+          source .venv/bin/activate
+          python -c "import jax; print(f'JAX version: {jax.__version__}')"
+          pytest -v --cov=zeropybench --cov-report=xml tests
 
       - name: Install system dependencies
         run: |
@@ -72,7 +75,9 @@ jobs:
 
       - name: Run notebook tests
         if: matrix.python.version == '3.14'
-        run: uv run pytest -v docs
+        run: |
+          source .venv/bin/activate
+          pytest -v docs
 
       - name: Upload coverage to Codecov
         if: matrix.python.version == '3.12'

pyproject.toml

@@ -49,7 +49,7 @@ docs = [
     'pillow>=9.0.0',
     'visu-hlo>=0.1.0',
 ]
-cuda12 = ['jax[cuda12]>=0.4']
+cuda12 = ['jax[cuda12]>=0.6']
 cuda13 = ['jax[cuda13]>=0.7; python_version>="3.11"']
 
 [tool.coverage.report]

src/zeropybench/_benchmark.py

@@ -213,7 +213,7 @@ def _compile_jax(self, globals: dict[str, Any]) -> tuple[str, float, bool]:
         compilation_time = time.perf_counter() - start_time
 
         hlo = compiled.as_text()
-        is_single_array = isinstance(compiled.out_info, jax.ShapeDtypeStruct)
+        is_single_array = isinstance(lowered.out_info, jax.ShapeDtypeStruct)
         return hlo, compilation_time, is_single_array
 
     def _run_many_times(