jaxdecomp 0.3.2__tar.gz → 0.3.4__tar.gz

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/.github/workflows/ci.yml

@@ -25,6 +25,7 @@ jobs:
   test:
     name: Run tests on Python ${{ matrix.python-version }}
     runs-on: ubuntu-latest
+    timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:
@@ -48,10 +49,13 @@ jobs:
       - name: Run tests
         # --doctest-modules imports jax at collection time, before conftest's session
         # fixture runs, so the device-count flags must be set in the environment here.
+        # JAXDECOMP_SKIP_PRIME_TESTS: the prime-sized FFTs deadlock XLA's CPU collective
+        # rendezvous on 4-core runners (thread starvation), see the conftest note.
         env:
           JAX_PLATFORM_NAME: cpu
           XLA_FLAGS: "--xla_force_host_platform_device_count=8"
-        run: uv run pytest -v
+          JAXDECOMP_SKIP_PRIME_TESTS: '1'
+        run: uv run pytest -v --timeout=600 --timeout-method=thread
 
   docs:
     name: Build documentation

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxdecomp
-Version: 0.3.2
+Version: 0.3.4
 Summary: JAX bindings for the cuDecomp library
 Author: Wassim Kabalan, Francois Lanusse
 License: MIT License
@@ -35,16 +35,17 @@ Classifier: Intended Audience :: Science/Research
 Project-URL: Homepage, https://github.com/DifferentiableUniverseInitiative/jaxDecomp
 Requires-Python: >=3.12
 Requires-Dist: jaxtyping>=0.2.0
-Requires-Dist: jax>=0.11.1
+Requires-Dist: jax>=0.11.2
 Provides-Extra: test
 Requires-Dist: pytest>=8.0.0; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: pytest-timeout; extra == "test"
 Requires-Dist: chex; extra == "test"
 Provides-Extra: benchmark
 Requires-Dist: jax-hpc-profiler>=0.3.2; extra == "benchmark"
 Provides-Extra: cuda
-Requires-Dist: jax[cuda12]>=0.11.1; extra == "cuda"
+Requires-Dist: jax[cuda12]>=0.11.2; extra == "cuda"
 Description-Content-Type: text/markdown
 
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 ]
 dependencies = [
   "jaxtyping>=0.2.0",
-  "jax>=0.11.1",
+  "jax>=0.11.2",
 ]
 
 [project.optional-dependencies]
@@ -39,13 +39,14 @@ test = [
   "pytest>=8.0.0",
   "pytest-cov",
   "pytest-mock",
+  "pytest-timeout",
   "chex",
 ]
 benchmark = [
   "jax-hpc-profiler>=0.3.2",
 ]
 cuda = [
-  "jax[cuda12]>=0.11.1",
+  "jax[cuda12]>=0.11.2",
 ]
 
 [dependency-groups]

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/cudecomp/fft.py

@@ -9,7 +9,7 @@ from jax import ShapeDtypeStruct
 from jax._src.interpreters import ad, mlir
 from jax._src.numpy.util import promote_dtypes_complex
 from jax.core import ShapedArray
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxdecomplib import _jaxdecomp
@@ -483,6 +483,8 @@ class PfftHiPrim(HiPrim):
         y_dot = self.expand(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self.expand(x), (self.params['fft_type'], self.params['adjoint'])
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/cudecomp/halo.py

@@ -9,7 +9,7 @@ from jax import ShapeDtypeStruct
 from jax._src.interpreters import ad, mlir
 from jax._src.typing import Array
 from jax.core import ShapedArray
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxdecomplib import _jaxdecomp
@@ -436,6 +436,8 @@ class HaloExchangeHiPrim(HiPrim):
         y_dot = self.expand(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self.expand(x), x
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/cudecomp/transpose.py

@@ -10,7 +10,7 @@ from jax._src.interpreters import ad, mlir
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.typing import Array, ArrayLike
 from jax.core import ShapedArray
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxdecomplib import _jaxdecomp
@@ -557,6 +557,8 @@ class TransposeHiPrim(HiPrim):
         y_dot = self.expand(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self.expand(x), None
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/jax/fft.py

@@ -6,7 +6,7 @@ from jax._src.api import ShapeDtypeStruct
 from jax._src.core import ShapedArray
 from jax._src.interpreters import batching
 from jax._src.typing import Array
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxdecomplib import _jaxdecomp
@@ -613,6 +613,8 @@ class PfftHiPrim(HiPrim):
         y_dot = self(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self(x), (self.params['fft_type'], self.params['adjoint'])
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/jax/halo.py

@@ -4,7 +4,7 @@ import jax
 from jax import ShapeDtypeStruct, lax
 from jax._src.interpreters import batching
 from jax.core import ShapedArray
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxdecomplib import _jaxdecomp
@@ -508,6 +508,8 @@ class HaloExchangeHiPrim(HiPrim):
         y_dot = self(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self(x), x
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/src/jaxdecomp/_src/jax/transpose.py

@@ -5,7 +5,7 @@ from jax import ShapeDtypeStruct, lax
 from jax._src.interpreters import batching
 from jax._src.typing import Array, ArrayLike
 from jax.core import ShapedArray
-from jax.experimental.hijax import VJPHiPrimitive as HiPrim
+from jax.experimental.hijax import HiPrim, linearize_from_jvp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 
@@ -360,6 +360,8 @@ class TransposeHiPrim(HiPrim):
         y_dot = self(x_dot)
         return y, y_dot
 
+    lin, linearized = linearize_from_jvp
+
     def vjp_fwd(self, nzs_in, x):
         return self(x), None
 

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/tests/conftest.py

@@ -6,6 +6,12 @@ import pytest
 setup_done = False
 on_cluster = False
 
+# XLA CPU collective starvation: on core-limited machines (e.g. 4-core CI runners), the
+# prime-sized shapes exercise long-running all_to_all collectives that deadlock the XLA CPU
+# rendezvous. CI sets JAXDECOMP_SKIP_PRIME_TESTS=1 to skip them; local machines and the
+# cluster (where the tests pass and matter) run them by default.
+skip_prime_tests = os.environ.get('JAXDECOMP_SKIP_PRIME_TESTS', '0') == '1'
+
 
 def is_on_cluster():
     global on_cluster

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/tests/test_fft.py

@@ -6,6 +6,7 @@ from conftest import (
     create_spmd_array,
     initialize_distributed,
     is_on_cluster,
+    skip_prime_tests,
 )
 
 initialize_distributed()
@@ -46,7 +47,13 @@ decomp = [
 global_shapes = [
     pytest.param((8, 16, 32), id='8x16x32'),
     pytest.param((8, 8, 8), id='8x8x8'),
-    pytest.param((29 * size, 19 * size, 17 * size), id='prime_sizes'),
+    pytest.param(
+        (29 * size, 19 * size, 17 * size),
+        id='prime_sizes',
+        marks=pytest.mark.skipif(
+            skip_prime_tests, reason='prime shapes deadlock XLA CPU collectives on core-limited runners (JAXDECOMP_SKIP_PRIME_TESTS=1)'
+        ),
+    ),
 ]  # Cubes, non-cubes and primes
 local_transpose = [
     pytest.param(True, id='local_transpose'),
@@ -309,6 +316,54 @@ class TestFFTsGrad:
         global_shape = (4, 4, 4)
         self._run_hessian_test(pdims, global_shape, local_transpose, 'cuDecomp', use_shardy, axis_type)
 
+    def _run_linearize_test(self, local_transpose, use_shardy, backend):
+        """Test that jax.linearize works with FFT (the HiPrim `lin`/`linearized` rules).
+
+        The linearized map must agree with the jvp rule (the FFT is linear) and satisfy
+        the transpose dot-product test ⟨w, J v⟩ == ⟨Jᵀ w, v⟩.
+        """
+        if use_shardy and not ALLOW_SHARDY_PARTITIONER:
+            pytest.skip(reason='Shardy partitioner is not supported in this JAX version use at least JAX 0.7.0')
+
+        jaxdecomp.config.update('transpose_axis_contiguous', local_transpose)
+        jax.config.update('jax_use_shardy_partitioner', use_shardy)
+
+        def f(arr):
+            return jaxdecomp.fft.pfft3d(arr, backend=backend)
+
+        # Small non-distributed array, like the hessian test
+        x = jax.random.normal(jax.random.PRNGKey(0), (4, 4, 4), dtype=jnp.float64)
+        v = jax.random.normal(jax.random.PRNGKey(1), x.shape, dtype=x.dtype)
+
+        # jax.linearize requires the HiPrim `lin`/`linearized` rules
+        _, lin_fn = jax.linearize(f, x)
+
+        # The linearized map must agree with the jvp rule (the FFT is linear)
+        _, Jv_jvp = jax.jvp(f, (x,), (v,))
+        Jv_lin = lin_fn(v)
+        assert assert_allclose(Jv_jvp, Jv_lin, rtol=1e-5, atol=1e-5)
+
+        # Transpose dot-product test: Re[wᵀ(J v)] == (Jᵀ w)ᵀv — the real part of the bilinear
+        # pairing, which is the transpose contract jax.linear_transpose defines for complex
+        # outputs on real inputs (the returned cotangent Jᵀ w is real).
+        w = jax.random.normal(jax.random.PRNGKey(2), Jv_lin.shape, dtype=Jv_lin.dtype)
+        _, vjp_fn = jax.vjp(f, x)
+        (JTw,) = vjp_fn(w)
+        assert assert_allclose(jnp.sum(w * Jv_lin).real, jnp.sum(JTw * v), rtol=1e-5, atol=1e-5)
+
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('local_transpose', local_transpose)
+    def test_jax_linearize(self, local_transpose, use_shardy):
+        """Test that jax.linearize works with FFT (JAX backend)."""
+        self._run_linearize_test(local_transpose, use_shardy, 'jax')
+
+    @pytest.mark.skipif(not is_on_cluster(), reason='Only run on cluster')
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('local_transpose', local_transpose)
+    def test_cudecomp_linearize(self, local_transpose, use_shardy):
+        """Test that jax.linearize works with FFT (cuDecomp backend)."""
+        self._run_linearize_test(local_transpose, use_shardy, 'cuDecomp')
+
     @pytest.mark.skipif(not is_on_cluster(), reason='Only run on cluster')
     @pytest.mark.parametrize('local_transpose', local_transpose)  # Test with and without local transpose
     @pytest.mark.parametrize('pdims', decomp)  # Test with Slab and Pencil decompositions

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/tests/test_halo.py

@@ -440,6 +440,66 @@ class TestHaloExchangeGrad:
     def test_cudecomp_hessian(self, pdims, use_shardy, axis_type):
         self._run_hessian_test(pdims, 'CUDECOMP', use_shardy, axis_type)
 
+    def _run_linearize_test(self, pdims, backend, use_shardy):
+        """Test that jax.linearize works with halo exchange (the HiPrim `lin`/`linearized` rules).
+
+        The linearized map must agree with the jvp rule (halo exchange is linear) and
+        satisfy the transpose dot-product test ⟨w, J v⟩ == ⟨Jᵀ w, v⟩.
+        """
+        jax.config.update('jax_use_shardy_partitioner', use_shardy)
+
+        if use_shardy and not ALLOW_SHARDY_PARTITIONER:
+            pytest.skip(reason='Shardy partitioner is not supported in this JAX version use at least JAX 0.7.0')
+
+        halo_size = 2
+
+        halo_x = (halo_size, halo_size) if pdims[0] > 1 else (0, 0)
+        halo_y = (halo_size, halo_size) if pdims[1] > 1 else (0, 0)
+        halo_extents = (halo_x[0], halo_y[0])
+        periodic = (True, True)
+        padding = (halo_x, halo_y, (0, 0))
+
+        def f(arr):
+            return jaxdecomp.halo_exchange(arr, halo_extents=halo_extents, halo_periods=periodic, backend=backend)
+
+        # Small non-distributed array, like the hessian test
+        padded = jnp.pad(
+            jax.random.normal(jax.random.PRNGKey(0), (16, 16, 16), dtype=jnp.float64),
+            padding,
+            mode='linear_ramp',
+            end_values=20,
+        )
+        v = jax.random.normal(jax.random.PRNGKey(1), padded.shape, dtype=padded.dtype)
+
+        # jax.linearize requires the HiPrim `lin`/`linearized` rules
+        _, lin_fn = jax.linearize(f, padded)
+
+        # The linearized map must agree with the jvp rule (halo exchange is linear)
+        _, Jv_jvp = jax.jvp(f, (padded,), (v,))
+        Jv_lin = lin_fn(v)
+        assert assert_allclose(Jv_jvp, Jv_lin, rtol=1e-5, atol=1e-5)
+
+        # Transpose dot-product test: Re[wᵀ(J v)] == (Jᵀ w)ᵀv — the real part of the bilinear
+        # pairing, which is the transpose contract jax.linear_transpose defines for complex
+        # outputs on real inputs (the returned cotangent Jᵀ w is real).
+        w = jax.random.normal(jax.random.PRNGKey(2), Jv_lin.shape, dtype=Jv_lin.dtype)
+        _, vjp_fn = jax.vjp(f, padded)
+        (JTw,) = vjp_fn(w)
+        assert assert_allclose(jnp.sum(w * Jv_lin).real, jnp.sum(JTw * v), rtol=1e-5, atol=1e-5)
+
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('pdims', pdims)
+    def test_jax_linearize(self, pdims, use_shardy):
+        """Test that jax.linearize works with halo exchange (JAX backend)."""
+        self._run_linearize_test(pdims, 'JAX', use_shardy)
+
+    @pytest.mark.skipif(not is_on_cluster(), reason='Only run on cluster')
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('pdims', pdims)
+    def test_cudecomp_linearize(self, pdims, use_shardy):
+        """Test that jax.linearize works with halo exchange (cuDecomp backend)."""
+        self._run_linearize_test(pdims, 'CUDECOMP', use_shardy)
+
     @pytest.mark.parametrize('use_shardy', use_shardy)  # Test with and without shardy
     @pytest.mark.parametrize('pdims', pdims)
     def test_jax_halo(

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/tests/test_transpose.py

@@ -8,6 +8,7 @@ from conftest import (
     create_spmd_array,
     initialize_distributed,
     is_on_cluster,
+    skip_prime_tests,
 )
 
 initialize_distributed()
@@ -45,7 +46,17 @@ pencil_1 = (size // 2, size // (size // 2))  # 2x2 for V100 and 4x2 for A100
 pencil_2 = (size // (size // 2), size // 2)  # 2x2 for V100 and 2x4 for A100
 
 decomp = [(size, 1), (1, size), pencil_1, pencil_2]
-global_shapes = [(8, 16, 32), (8, 8, 8), (29 * size, 19 * size, 17 * size)]
+global_shapes = [
+    pytest.param((8, 16, 32), id='8x16x32'),
+    pytest.param((8, 8, 8), id='8x8x8'),
+    pytest.param(
+        (29 * size, 19 * size, 17 * size),
+        id='prime_sizes',
+        marks=pytest.mark.skipif(
+            skip_prime_tests, reason='prime shapes deadlock XLA CPU collectives on core-limited runners (JAXDECOMP_SKIP_PRIME_TESTS=1)'
+        ),
+    ),
+]
 # Cubes, non-cubes and primes
 local_transpose = [
     pytest.param(False, id='no_local_transpose'),
@@ -314,6 +325,54 @@ class TestTransposesGrad:
         global_shape = (4, 4, 4)
         self._run_hessian_test(pdims, global_shape, local_transpose, 'cuDecomp', use_shardy, axis_type)
 
+    def _run_linearize_test(self, local_transpose, use_shardy, backend):
+        """Test that jax.linearize works with transpose (the HiPrim `lin`/`linearized` rules).
+
+        The linearized map must agree with the jvp rule (the transpose is linear) and
+        satisfy the transpose dot-product test ⟨w, J v⟩ == ⟨Jᵀ w, v⟩.
+        """
+        if use_shardy and not ALLOW_SHARDY_PARTITIONER:
+            pytest.skip(reason='Shardy partitioner is not supported in this JAX version use at least JAX 0.7.0')
+
+        jaxdecomp.config.update('transpose_axis_contiguous', local_transpose)
+        jax.config.update('jax_use_shardy_partitioner', use_shardy)
+
+        def f(arr):
+            return transposeXtoY(arr, backend=backend)
+
+        # Small non-distributed array, like the hessian test
+        x = jax.random.normal(jax.random.PRNGKey(0), (4, 4, 4), dtype=jnp.float64)
+        v = jax.random.normal(jax.random.PRNGKey(1), x.shape, dtype=x.dtype)
+
+        # jax.linearize requires the HiPrim `lin`/`linearized` rules
+        _, lin_fn = jax.linearize(f, x)
+
+        # The linearized map must agree with the jvp rule (the transpose is linear)
+        _, Jv_jvp = jax.jvp(f, (x,), (v,))
+        Jv_lin = lin_fn(v)
+        assert assert_allclose(Jv_jvp, Jv_lin, rtol=1e-5, atol=1e-5)
+
+        # Transpose dot-product test: Re[wᵀ(J v)] == (Jᵀ w)ᵀv — the real part of the bilinear
+        # pairing, which is the transpose contract jax.linear_transpose defines for complex
+        # outputs on real inputs (the returned cotangent Jᵀ w is real).
+        w = jax.random.normal(jax.random.PRNGKey(2), Jv_lin.shape, dtype=Jv_lin.dtype)
+        _, vjp_fn = jax.vjp(f, x)
+        (JTw,) = vjp_fn(w)
+        assert assert_allclose(jnp.sum(w * Jv_lin).real, jnp.sum(JTw * v), rtol=1e-5, atol=1e-5)
+
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('local_transpose', local_transpose)
+    def test_jax_linearize(self, local_transpose, use_shardy):
+        """Test that jax.linearize works with transpose (JAX backend)."""
+        self._run_linearize_test(local_transpose, use_shardy, 'jax')
+
+    @pytest.mark.skipif(not is_on_cluster(), reason='Only run on cluster')
+    @pytest.mark.parametrize('use_shardy', use_shardy)
+    @pytest.mark.parametrize('local_transpose', local_transpose)
+    def test_cudecomp_linearize(self, local_transpose, use_shardy):
+        """Test that jax.linearize works with transpose (cuDecomp backend)."""
+        self._run_linearize_test(local_transpose, use_shardy, 'cuDecomp')
+
     @pytest.mark.skipif(not is_on_cluster(), reason='Only run on cluster')
     # Cartesian product test
     @pytest.mark.parametrize('pdims', decomp)

{jaxdecomp-0.3.2 → jaxdecomp-0.3.4}/uv.lock

@@ -592,7 +592,7 @@ wheels = [
 
 [[package]]
 name = "jax"
-version = "0.11.1"
+version = "0.11.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jaxlib" },
@@ -601,9 +601,9 @@ dependencies = [
     { name = "opt-einsum" },
     { name = "scipy" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/90/aa/c3b6ac3863a490661e255d2356802b54d57b533460fa221d30743d81a6c1/jax-0.11.1.tar.gz", hash = "sha256:43e0d45b1dac002eca04c2de73298395225dd0d4651e3f573a540eaa18abfd80", size = 2864119, upload-time = "2026-08-17T20:31:29.938Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/96/c3/cbb70e5e0846891b45c9b02b0755235618043cd2a8edd1f5bfe24b077cf4/jax-0.11.2.tar.gz", hash = "sha256:540dc0bed96bd5a0d8acca15cf16198fea98ef2fc151d40bdf1ace49eb76efd9", size = 2904036, upload-time = "2026-09-17T23:43:35.896Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8f/f4/5f9a286e1eeca9faa8b83885edccf63ab2f1b22e17f6a5ea66ac86b1afe4/jax-0.11.1-py3-none-any.whl", hash = "sha256:72ed60bf85a0b53d0f5b5cd61e37a8b25f369f6f88481feb98ab489894861a82", size = 3302513, upload-time = "2026-08-17T20:29:00.294Z" },
+    { url = "https://files.pythonhosted.org/packages/ca/87/0c0ca8433e1135f6d1acd355cff21f31c3d433d937d6547aa545c8b060b8/jax-0.11.2-py3-none-any.whl", hash = "sha256:59e7ed9bd9ace049f4856752fe580c890e0164e1c023f40fadac386408daad66", size = 3345045, upload-time = "2026-09-17T23:41:14.517Z" },
 ]
 
 [package.optional-dependencies]
@@ -614,33 +614,33 @@ cuda12 = [
 
 [[package]]
 name = "jax-cuda12-pjrt"
-version = "0.11.1"
+version = "0.11.2"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/0e/3e/d2072547b8738a7c4cf9ad79589d0efed1a4700952be30b1faa88f87969d/jax_cuda12_pjrt-0.11.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:40021d455bd94522704e321d30e8ada26a7c0a6ef95c4a7542abda0138a29d1c", size = 179725961, upload-time = "2026-08-17T20:29:04.18Z" },
-    { url = "https://files.pythonhosted.org/packages/55/15/5709e379c1191b1d354e74a04cb01d61abe29ce4bba29833229525882566/jax_cuda12_pjrt-0.11.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:fdd9d72408e17f7a23bdfa266e9c90136620a724df584629f3506722e2b2f434", size = 184904180, upload-time = "2026-08-17T20:29:12.991Z" },
+    { url = "https://files.pythonhosted.org/packages/73/a8/03be433dec7574ece4d45b90d709e25ee0b9e19713e73fbfc52d6384de75/jax_cuda12_pjrt-0.11.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:f36247d40973b20acf3c69dac7c9c720daaa4096a809b3e8fd70292be617cdeb", size = 182867418, upload-time = "2026-09-17T23:41:18.187Z" },
+    { url = "https://files.pythonhosted.org/packages/38/57/2df730cd97270f212938f1b86ba50d6118220fb9e4c2dda08d260418679e/jax_cuda12_pjrt-0.11.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:2133b3afc63307c016ae4980f166b2b3de199759363ead74f7d550eaae953463", size = 188235313, upload-time = "2026-09-17T23:41:24.135Z" },
 ]
 
 [[package]]
 name = "jax-cuda12-plugin"
-version = "0.11.1"
+version = "0.11.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jax-cuda12-pjrt" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b6/55/7d99e5155d8c1a402c3fcd8f5e3ac9d7d3123725d1cef9d7f519eb4408f4/jax_cuda12_plugin-0.11.1-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:f2a382e5dbf66330740bf4db3903c612e84b24de09913c7953ed2403c296b686", size = 8551003, upload-time = "2026-08-17T20:29:16.516Z" },
-    { url = "https://files.pythonhosted.org/packages/56/45/e30b3348126ded20d5b388e18a4fe399259538506ba1abc2812913f21eee/jax_cuda12_plugin-0.11.1-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:f0154fbf236bb2231a6fca5204dd1e43ac0c1cf0de3a0020fff9103e53fb8572", size = 8594678, upload-time = "2026-08-17T20:29:18.489Z" },
-    { url = "https://files.pythonhosted.org/packages/fb/28/8291d90302117e788cc40f10e4cb442444ac8bf21673ff73ef9c21f8243a/jax_cuda12_plugin-0.11.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:e8a8d66f43d00b35c0430366b117c1b61283764f0d10d3965abfe39f8728486f", size = 8550626, upload-time = "2026-08-17T20:29:20.353Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/e7/d2407b86fcb882ba0ce16a6ff23d0f21d18e3725133d8ac0df7497d704a8/jax_cuda12_plugin-0.11.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:d8efc51fd4998a2990c467e63c386cfdc9187c7d105825c5c4ab743e7d522d11", size = 8594940, upload-time = "2026-08-17T20:29:22.337Z" },
-    { url = "https://files.pythonhosted.org/packages/fb/e0/de583e82a8979741ddaa8116c7d151b75a4491ccb526489d4c8043630411/jax_cuda12_plugin-0.11.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:3ca119b301d9a93bcb8375c6937cf79492ff5951ca5cbb8b46f7e34af42dcfd3", size = 8551250, upload-time = "2026-08-17T20:29:24.075Z" },
-    { url = "https://files.pythonhosted.org/packages/bc/f1/edda4a57aad1f6975bafe3abdae3297fe4cf52defe2ffe167e58093a1bcb/jax_cuda12_plugin-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:6bb8f106821517e34488a5320e45f5def4a4628f2a254afa6e2c7eec35419bbe", size = 8595070, upload-time = "2026-08-17T20:29:25.868Z" },
-    { url = "https://files.pythonhosted.org/packages/ca/c0/b880ab87b07d3b7bce14fdb965cbe9fb6b3a1b61d73b1d6886da01a3d293/jax_cuda12_plugin-0.11.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:479407a4eed9fb5b1fad37ae3a6a1c9e8284bdb2e52ba6dfc0e22ac595bb94e8", size = 8565526, upload-time = "2026-08-17T20:29:27.915Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/9f/ce21715c1788ca6686c0550feeef28d06f574872e4619a74ac729b8ff47d/jax_cuda12_plugin-0.11.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:771a43862a86874b114c0d55f23136c8a228e3d47e7718dc4ed60de84ac36d89", size = 8605791, upload-time = "2026-08-17T20:29:29.551Z" },
-    { url = "https://files.pythonhosted.org/packages/49/f3/e570a3c7a77d8b36317b214460515f4332b0960efeaafed8c00f46d2d646/jax_cuda12_plugin-0.11.1-cp315-cp315-manylinux_2_27_aarch64.whl", hash = "sha256:4a09e11b9c2704c6af53bcd1572066429b015a36a106882641e127f30e7a4589", size = 8551262, upload-time = "2026-08-17T20:29:31.169Z" },
-    { url = "https://files.pythonhosted.org/packages/4c/eb/c521edf23104635e79ecba0c5881e86b85eaf632241d911220a64ba926b5/jax_cuda12_plugin-0.11.1-cp315-cp315-manylinux_2_27_x86_64.whl", hash = "sha256:d87516db355e761eb98efefd817a5e985688bad32d3c05a93eb7b64ed721b3d8", size = 8595108, upload-time = "2026-08-17T20:29:33.159Z" },
-    { url = "https://files.pythonhosted.org/packages/b0/81/ae93fb9edf4fdbace632ce1e649a3f8965311efd40e22d7bf0c51c2d5246/jax_cuda12_plugin-0.11.1-cp315-cp315t-manylinux_2_27_aarch64.whl", hash = "sha256:b6c88d844e26444e78dc05d8d15342d84e15aca79ab34881020e49f252fa8024", size = 8565490, upload-time = "2026-08-17T20:29:34.848Z" },
-    { url = "https://files.pythonhosted.org/packages/68/ec/0f85c6d5eec3f2f83ceda1c6d8ab45a7dc025914cd50dd4d82b75e483510/jax_cuda12_plugin-0.11.1-cp315-cp315t-manylinux_2_27_x86_64.whl", hash = "sha256:055d8835b603f5a7a374de5342bb5cfcb0a4bbbb8525e940bbc51aee383ecdbe", size = 8606031, upload-time = "2026-08-17T20:29:37.016Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/e9/f205c95df880618d5691ad5444e8798368a3cf76f1380e153f78b392470e/jax_cuda12_plugin-0.11.2-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:ad0fbb60548585d5f0df2fe11eb8e6d41a9856274ed9bc3ebe2d1e1e1c3d012f", size = 8603910, upload-time = "2026-09-17T23:41:27.572Z" },
+    { url = "https://files.pythonhosted.org/packages/77/db/ca022a589ff15628827498dcedeae72fa70396f42b58469a9de4425afc99/jax_cuda12_plugin-0.11.2-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:ae67a4cc49b34c575713a8311f7597f6700ec5a2b6afbdbc3db5550fb022873c", size = 8636558, upload-time = "2026-09-17T23:41:29.657Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/13/a7d6099008e001bf71d3273d21fe970ee768cab3600e31f49ae82b578d11/jax_cuda12_plugin-0.11.2-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:291aadea4a52c53a35f957ca5a3c0dd5316b5c082fcbb51f164b7f20d4e1d659", size = 8603926, upload-time = "2026-09-17T23:41:31.82Z" },
+    { url = "https://files.pythonhosted.org/packages/57/ba/26bd500531d81b16d66e4a90bb636fbc0e9e6f921b47a3169007159f694c/jax_cuda12_plugin-0.11.2-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:55ed6dbdf4d6f57c6c7b0aadaeb8e337371a25249a0e3c1ef9efcb607b0bbc67", size = 8636734, upload-time = "2026-09-17T23:41:33.538Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/56/ed9f440887dae68fb3adc6f011c1cdeeb05774f2d795ff6c64eb4dbb95c6/jax_cuda12_plugin-0.11.2-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:ec7e06d7a450807967f35091d4af6c0a72461bdc77c4b21718a50d8377eaeafb", size = 8603321, upload-time = "2026-09-17T23:41:35.155Z" },
+    { url = "https://files.pythonhosted.org/packages/02/14/fcae3681051de6f109c85a9199220e4f4aafa4698acbe131084e56f790e9/jax_cuda12_plugin-0.11.2-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:96d198157e76ead91958742466f5c7462618dd27831dc82384880e2e80ebe4e0", size = 8636844, upload-time = "2026-09-17T23:41:37.232Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/bb/64db10a05cd2821b6b8c86896fdc461c13bdc5dba10d19bcf271d7327c00/jax_cuda12_plugin-0.11.2-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:af5147dde561d8169a5073fe64aa75763b527d65a7087fe9774cb80be84fdae6", size = 8619610, upload-time = "2026-09-17T23:41:39.435Z" },
+    { url = "https://files.pythonhosted.org/packages/06/ca/67b5934f7d7e162dd6389a7c23aaa81b08991ae51d392f9153d34eb9f735/jax_cuda12_plugin-0.11.2-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:c0a3920ae96cd439557ab8ba510b2fdbad0c935c1d95cf5bb54113e025ba9a82", size = 8652035, upload-time = "2026-09-17T23:41:41.069Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/fc/4bab2b740004e3f449c38c44d6c48291708582051517ffe97865220960b9/jax_cuda12_plugin-0.11.2-cp315-cp315-manylinux_2_27_aarch64.whl", hash = "sha256:a2565efe1dad968fe18e04e39161f059437d923f8bed966877b33721ccbab5d5", size = 8606679, upload-time = "2026-09-17T23:41:42.902Z" },
+    { url = "https://files.pythonhosted.org/packages/86/8b/ac014271a8e01f35891ef57c3788fcb8ecd00f75f8638594af638510d9bc/jax_cuda12_plugin-0.11.2-cp315-cp315-manylinux_2_27_x86_64.whl", hash = "sha256:8044b97ca41ecc835034c0f780abdf79dd1d673834031a2710882f895429679a", size = 8640048, upload-time = "2026-09-17T23:41:45.15Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/8d/c816e5c281aa16956b88838901a1cd47a99760df45113d79507aecb41489/jax_cuda12_plugin-0.11.2-cp315-cp315t-manylinux_2_27_aarch64.whl", hash = "sha256:614eb748b5175323cd5865653a642d1b7e1af69f3afa47565f8fddb21ad011fd", size = 8621068, upload-time = "2026-09-17T23:41:46.792Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/c1/e51e5225a583cf9e10565978be95e87dda09c226bd41372df6c251cecb24/jax_cuda12_plugin-0.11.2-cp315-cp315t-manylinux_2_27_x86_64.whl", hash = "sha256:e6e154479252c1fbe7dccbb3be5fd90b5f91537786dde4979b43acac4c84a07d", size = 8654513, upload-time = "2026-09-17T23:41:48.65Z" },
 ]
 
 [package.optional-dependencies]
@@ -698,6 +698,7 @@ test = [
     { name = "pytest" },
     { name = "pytest-cov" },
     { name = "pytest-mock" },
+    { name = "pytest-timeout" },
 ]
 
 [package.dev-dependencies]
@@ -718,13 +719,14 @@ docs = [
 [package.metadata]
 requires-dist = [
     { name = "chex", marker = "extra == 'test'" },
-    { name = "jax", specifier = ">=0.11.1" },
-    { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda'", specifier = ">=0.11.1" },
+    { name = "jax", specifier = ">=0.11.2" },
+    { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda'", specifier = ">=0.11.2" },
     { name = "jax-hpc-profiler", marker = "extra == 'benchmark'", specifier = ">=0.3.2" },
     { name = "jaxtyping", specifier = ">=0.2.0" },
     { name = "pytest", marker = "extra == 'test'", specifier = ">=8.0.0" },
     { name = "pytest-cov", marker = "extra == 'test'" },
     { name = "pytest-mock", marker = "extra == 'test'" },
+    { name = "pytest-timeout", marker = "extra == 'test'" },
 ]
 provides-extras = ["test", "benchmark", "cuda"]
 
@@ -745,7 +747,7 @@ docs = [
 
 [[package]]
 name = "jaxlib"
-version = "0.11.1"
+version = "0.11.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "ml-dtypes" },
@@ -753,28 +755,28 @@ dependencies = [
     { name = "scipy" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/6a/b4/9a7caf8704536e694ab8b4286eb4afa2224f3f75283f370d2246343e7bbc/jaxlib-0.11.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:384f633331c012491ccd65c38d47cd80d4131f5b3852694766bffa9385ddc9f3", size = 63173264, upload-time = "2026-08-17T20:30:07.721Z" },
-    { url = "https://files.pythonhosted.org/packages/77/d9/26613efbc0e219ef16b6069c38378a9b360f85926573e1995c391e48d73e/jaxlib-0.11.1-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:47bab3a7da8a818f9e7ddecd4ed815429caa2840eda9e968bb6d33b9d8c52067", size = 82772925, upload-time = "2026-08-17T20:30:11.388Z" },
-    { url = "https://files.pythonhosted.org/packages/05/41/07c259ea10a5eeb100a66222507e2ce3a01474fa8f13c250bf1a3559a5bb/jaxlib-0.11.1-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:c3dcec3bb33ebf5c76dc33f962c796652c077edb246464fa32cb1d8fcb6087e6", size = 87903454, upload-time = "2026-08-17T20:30:14.889Z" },
-    { url = "https://files.pythonhosted.org/packages/77/d7/555041e5d4116c0aaf8ff513b380da38df547844c30c7c1d41c36e18d12b/jaxlib-0.11.1-cp312-cp312-win_amd64.whl", hash = "sha256:eef943b42d6d42a2b2f3dc2a66e7a3dc89aeb1541f5a13ad8e1e94001bb590ec", size = 68531585, upload-time = "2026-08-17T20:30:18.488Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/22/d076bf545853d1f60ef906cb6d69751f9b149f857f4390de989d5165ed80/jaxlib-0.11.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c7518af3168ab4cf0a00b92f5463018baf0dad3faa3b358a3d688996afa6fd50", size = 63174384, upload-time = "2026-08-17T20:30:21.949Z" },
-    { url = "https://files.pythonhosted.org/packages/a8/3c/5b0ce63c880d4f9cef492f04e11be7a60bdd3f529311839652b928e8b320/jaxlib-0.11.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:3a679a49b8d329b6e431f3aeb666467c6efac750e3b2ee90d0e74efe7959e643", size = 82774089, upload-time = "2026-08-17T20:30:25.232Z" },
-    { url = "https://files.pythonhosted.org/packages/f5/28/ef1371618e784880e68c2ab5277d5ccbaf518a2a5ec9ec05ef90a9ee7d1b/jaxlib-0.11.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:7ebe86a7b891fcda83e7f634a677d285681a7b80b3ec9ed435536078a8371acf", size = 87901562, upload-time = "2026-08-17T20:30:28.873Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/c7/f92d23dc89bdbe5193ad139aa8973eab485898f330de2033bc4fdd8ac4c9/jaxlib-0.11.1-cp313-cp313-win_amd64.whl", hash = "sha256:8379a8902db216f21673e7591acedc8f925658c59d9d2ec9afd5e2244246986a", size = 68531205, upload-time = "2026-08-17T20:30:32.376Z" },
-    { url = "https://files.pythonhosted.org/packages/a8/90/29f37540af0d463838eaac236d58267815a972a5e2f5e59dc6e1992c5dbc/jaxlib-0.11.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:a063d6110d296a935933c35aa7ed61b3a22ae2466e38834eecd48ae2d14fb8a3", size = 63174889, upload-time = "2026-08-17T20:30:35.472Z" },
-    { url = "https://files.pythonhosted.org/packages/35/7f/263b873ed570567880e3ec2acb13eb5cf184972b5b79d883dbccc7d58c4a/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:226616df027f1348da77764753c6c1e2f4d45b5551b67bcfdcafcf9cadf2e280", size = 82773791, upload-time = "2026-08-17T20:30:38.783Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/a2/07dba61656a7fd3b0db734834158e9abcc56372592d9326cf39a72c36762/jaxlib-0.11.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:739ead792ddb854cb3ca5decfe00aab08a731131ab2cbeed9b74f4545c92fd01", size = 87903999, upload-time = "2026-08-17T20:30:42.645Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/4b/a925e73d19a7cb60c25396db208a1dd43ab7ae9f77a85646e65c02c3e961/jaxlib-0.11.1-cp314-cp314-win_amd64.whl", hash = "sha256:3c2788a45d334be1240547a31df0f8af5f1b59e95984c3bc86d9badce9293685", size = 71068954, upload-time = "2026-08-17T20:30:47.013Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/2f/d4eef297c2f8c3dd5e2ba1afb12d14328c13d343681d70d6d5eb95e53f17/jaxlib-0.11.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8dfbca728bd534c0690514527a15db52a8e74f7bae406984b57a3c0c99e75ab6", size = 63263304, upload-time = "2026-08-17T20:30:50.17Z" },
-    { url = "https://files.pythonhosted.org/packages/60/83/0eb0b2b7d9a32a742e351bf8f94663b3759408c579bd8964628fe9e0bcbe/jaxlib-0.11.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:f9fc4eccfb882eab746ac75862e82833eab1a72d864dcb20bdc83cdf295d81bd", size = 82866887, upload-time = "2026-08-17T20:30:53.518Z" },
-    { url = "https://files.pythonhosted.org/packages/d5/c6/d2e03edf89d577241aaa70a82a21a7e2474a113af4f7e5b2bccd65f72a2e/jaxlib-0.11.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:12803e2d6ca821d9f59aeb31ceaa6e49403aaeefbb62e3cba90e1ebcbe244209", size = 87999113, upload-time = "2026-08-17T20:30:57.474Z" },
-    { url = "https://files.pythonhosted.org/packages/69/93/da40546a2f67735717a1ffb4751da4eb92527fa326caa6e9973444b41321/jaxlib-0.11.1-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:6e264bd7cecff9053f71c4c7d1aeb0940aa7da9226ddcb00117abb39b5faccf4", size = 63174333, upload-time = "2026-08-17T20:31:00.789Z" },
-    { url = "https://files.pythonhosted.org/packages/d9/da/d73f62bc74692640e0d140db8964fa98a403686eb22c8d0810776e3a2e9e/jaxlib-0.11.1-cp315-cp315-manylinux_2_27_aarch64.whl", hash = "sha256:93c945387919ae27e1475b86802f72ec9f511f65ec74677b33ba47f895e51f9c", size = 82774506, upload-time = "2026-08-17T20:31:03.949Z" },
-    { url = "https://files.pythonhosted.org/packages/d9/c4/44e62769133051f37e12e33cf41e50c24c229c9378624c076edf53f8dd73/jaxlib-0.11.1-cp315-cp315-manylinux_2_27_x86_64.whl", hash = "sha256:b953f341c67d28f202b8af6dfaf9908281e2129ad9bc100454a9fdd2dae0b6f2", size = 87904049, upload-time = "2026-08-17T20:31:12.41Z" },
-    { url = "https://files.pythonhosted.org/packages/75/d1/1243129103b5df137cc84d2bd998d403b5dd2ebcb857f51fd318ebd20282/jaxlib-0.11.1-cp315-cp315-win_amd64.whl", hash = "sha256:8f7956e26fc27b3782b180688545df6ac98c60b9429c45467eaa268bdc60f091", size = 71068945, upload-time = "2026-08-17T20:31:17.332Z" },
-    { url = "https://files.pythonhosted.org/packages/8f/17/bbc0a38d2d560c92e1dd123d293ec2f78e24a843baa13238fd361e7ab370/jaxlib-0.11.1-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:d3945259807c26d231a2b61a63f4b6f1058f6519803b9bcc55ea2f1ae6a13108", size = 63263361, upload-time = "2026-08-17T20:31:20.369Z" },
-    { url = "https://files.pythonhosted.org/packages/97/5b/819e20e911ba4eae21f96ec808a97605383d519d605276914a2b0fb8ef2a/jaxlib-0.11.1-cp315-cp315t-manylinux_2_27_aarch64.whl", hash = "sha256:3bce7f70714ed7e1967b9725e67c0aa3e111ca7e0dc7434c5d7d1e2047a8de26", size = 82868605, upload-time = "2026-08-17T20:31:23.805Z" },
-    { url = "https://files.pythonhosted.org/packages/8c/e7/9593d22d398b66dae5fab17f961ed16ce7acfe11425276a6d2a0200d97ff/jaxlib-0.11.1-cp315-cp315t-manylinux_2_27_x86_64.whl", hash = "sha256:a3064a14d47f03b777454667bc8bf0a7c2d506714fe108adacf734f63cd7a5ae", size = 87998599, upload-time = "2026-08-17T20:31:27.376Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/af/56ca13ef7a98137b087ef0c24dd5ace3163895f0164670babaa7dfcb69e6/jaxlib-0.11.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:eca80e745b7623369a76b5f685676ec233a33a55d69f60fa7e57cbefcfd47adc", size = 64867740, upload-time = "2026-09-17T23:42:19.368Z" },
+    { url = "https://files.pythonhosted.org/packages/69/72/318270378f5ff554ee921944bd97f59c0d6b35a8397cd35771dfdce1f445/jaxlib-0.11.2-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:d497004e80daa91b98258d645b374e8dff480bf017d2c5f3e30b1371ac8f54b1", size = 84659774, upload-time = "2026-09-17T23:42:22.846Z" },
+    { url = "https://files.pythonhosted.org/packages/51/d2/07c3ea737d65b328b00377462e0e5dc8993a8a88d2a8a89e1c42038a4e53/jaxlib-0.11.2-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:6338a1f0956f60d02529404ac33a8070015a52635cdb6bbb15d0d199757077bb", size = 89896206, upload-time = "2026-09-17T23:42:26.851Z" },
+    { url = "https://files.pythonhosted.org/packages/74/f7/8f6f7c55b80ca45e2225c78f4f44dbda8e5f420f933a5896e21045c3fb41/jaxlib-0.11.2-cp312-cp312-win_amd64.whl", hash = "sha256:53217a47d14eedd8fb5101f16aab1dfd58067835b4fb11910b50306218c7e620", size = 73563902, upload-time = "2026-09-17T23:42:30.216Z" },
+    { url = "https://files.pythonhosted.org/packages/47/7f/f0d33a414e2d84fca39438a30f0588c32f7365746e903293cadad9d52b0d/jaxlib-0.11.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:766bd90e27a0ff53b87bb9d70c73565ec80b9b1165320ccb8f72de46b85510df", size = 64868496, upload-time = "2026-09-17T23:42:33.478Z" },
+    { url = "https://files.pythonhosted.org/packages/67/9e/6f6307527079e4fa65af9a813c3f07ba598f57f0936045ea3b02cc9d2e01/jaxlib-0.11.2-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:5f9d3833cb2bea5e346bd2138ba74992017e0841c63d890a0157a13f9cbf439e", size = 84657374, upload-time = "2026-09-17T23:42:36.804Z" },
+    { url = "https://files.pythonhosted.org/packages/24/ec/014b428c3a05837643768d08aeb8fb1289d119d3f9f6fba479dffae05dbc/jaxlib-0.11.2-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:5f3cac8d7030c1f80a182025d7171ea3b1245162a710f6e58d25ee3fe1748aac", size = 89896413, upload-time = "2026-09-17T23:42:40.466Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/2f/72dfa1d0340866ecce9e111aac05a826ff116852402fa15c0db1b91b2d1b/jaxlib-0.11.2-cp313-cp313-win_amd64.whl", hash = "sha256:523e72d6e8d188b30bee9b6c62e5e9b769ac118aab5f07343588ba64d45cf9f3", size = 73564682, upload-time = "2026-09-17T23:42:44.134Z" },
+    { url = "https://files.pythonhosted.org/packages/fb/be/b5562a9ffab85bac6303f9d33af49c4bc51f85e33f406616d018567749fb/jaxlib-0.11.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:edad6b2a18d63dc2964cc97a16d126dad41621edac79d651e0f6c8785d891369", size = 64874852, upload-time = "2026-09-17T23:42:47.578Z" },
+    { url = "https://files.pythonhosted.org/packages/36/ca/5a5f933b5d6a281b8b758f45871e3b72fd2290dbd5c6946b93140dd24012/jaxlib-0.11.2-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:061468eb5ac6b6213215fd1102b237bcbc93f24480c881245b73619d0c2a321e", size = 84664654, upload-time = "2026-09-17T23:42:50.815Z" },
+    { url = "https://files.pythonhosted.org/packages/35/2e/85156ae0a42292ef8e648c503437c65e434f6c4370f44e7fe5d4f845e6da/jaxlib-0.11.2-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:1f31d8a6b1fb13181ff1807757506d10f6e7cba6b2707f7f80f73b853cb2ff22", size = 89901153, upload-time = "2026-09-17T23:42:54.265Z" },
+    { url = "https://files.pythonhosted.org/packages/c0/8a/2c8d8cfd1450d4fc209273d09440b771608f4f6e657829088aabcb57072e/jaxlib-0.11.2-cp314-cp314-win_amd64.whl", hash = "sha256:ea5e58673775ec7d3f46ade4e2d5a0d1673f30bb76f2a70c9b0036eafe5e4053", size = 76273242, upload-time = "2026-09-17T23:42:57.836Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/57/a19ec0aa17be6acc405b09c52a6d0454eabec6a8c4a1a4bacfbce747c2f4/jaxlib-0.11.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a17074e2c58fa5610f6ec61955939d53e6372f59931d75eb1010926e6d81565d", size = 64969331, upload-time = "2026-09-17T23:43:01.079Z" },
+    { url = "https://files.pythonhosted.org/packages/0a/c1/dd8d109f48b27f884b3205f805066982f2f1ac427f0f500d31b353555626/jaxlib-0.11.2-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:5b052113f7a8acb0b6f14b356c83bb0f1e7ba90df0b22bbcf476d3750de5ea22", size = 84768035, upload-time = "2026-09-17T23:43:04.511Z" },
+    { url = "https://files.pythonhosted.org/packages/a5/31/5876f836644db19abdb9a620111e3e6307165934d73fa3f32de9fcf62343/jaxlib-0.11.2-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:d100ddb54a68d59cc3103729082324fccc4ca65e1d1b03c325208c43dd5c3c40", size = 89991692, upload-time = "2026-09-17T23:43:08.438Z" },
+    { url = "https://files.pythonhosted.org/packages/bd/2f/3cd215083ed42ae60bef985b36255e16c07cc5b6311f88d107726518675b/jaxlib-0.11.2-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:c2563f58a5d16f6b8d68bb65c1ba83e9ec81bea0d93cd3b1a73b0a5ffe76c1bb", size = 64877295, upload-time = "2026-09-17T23:43:12.412Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/69/817015b0d00ab42efe7c6dad778c4c1ca5972b3e180ed27448adc620035b/jaxlib-0.11.2-cp315-cp315-manylinux_2_27_aarch64.whl", hash = "sha256:8df25d031827658fa434e6a8ff5c749700711fcc1553757d4bde4621019180cb", size = 84671133, upload-time = "2026-09-17T23:43:15.701Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/9a/6e8b7319a325d23fcd76442370a7ea2bf442858a0c8196ec4caf69d67344/jaxlib-0.11.2-cp315-cp315-manylinux_2_27_x86_64.whl", hash = "sha256:b2965e43db9951a4b5e11dfd437477944fa5a4cc4f4fb969725788b90ebd7c09", size = 89903001, upload-time = "2026-09-17T23:43:19.354Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/ee/6e653bbe8deb7303eb07b78fd36c9b57fd685641a62ab79532b885624939/jaxlib-0.11.2-cp315-cp315-win_amd64.whl", hash = "sha256:75db0922c72c511c71f9d022d7dcaca0d5492d30ad815041857fc447948c91f3", size = 76274240, upload-time = "2026-09-17T23:43:23.037Z" },
+    { url = "https://files.pythonhosted.org/packages/b2/6b/4b2f037c70a08d76786f2548b107c2a355f89bdd0c1fbba821b8e4776f15/jaxlib-0.11.2-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:dc154e762e6a5b19d5061e69397edccfea8c3acfbc831e487e2f6e403a218d3a", size = 64975259, upload-time = "2026-09-17T23:43:26.262Z" },
+    { url = "https://files.pythonhosted.org/packages/ef/8b/ee9ad489ae420f44f87c953ace859d7cb72dcc2ebd3112c50eae7b8beef5/jaxlib-0.11.2-cp315-cp315t-manylinux_2_27_aarch64.whl", hash = "sha256:33d9c43f729f0e5b597f3fca73b2e1b342877e063a7e259ab57b6cc6e686369e", size = 84759864, upload-time = "2026-09-17T23:43:29.829Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/a3/1f0b704e4b1c3a14bbc547b252e4e5b0f2d8082df1e8b664f8b579adabec/jaxlib-0.11.2-cp315-cp315t-manylinux_2_27_x86_64.whl", hash = "sha256:4b553584279c42018dae4ede7963c7db6cdb35d1f36804b18a07256146c4529f", size = 89995359, upload-time = "2026-09-17T23:43:33.461Z" },
 ]
 
 [[package]]
@@ -1371,7 +1373,7 @@ name = "nvidia-cublas-cu12"
 version = "12.9.2.10"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cuda-nvrtc-cu12" },
+    { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/a2/c96163a0fff1839c0c9548bbdeae7b853b867009e33b9b9264adc238b1cf/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5572131a59c3eebeeb1c4c8144f772d49372c20124916e072a0e3fc30df421d5", size = 575012079, upload-time = "2026-04-08T18:51:47.303Z" },
@@ -1428,7 +1430,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.25.1.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12" },
+    { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/9f/97/65cee6620521af847fde280f3bd10784f3ab3d513f1bfdd87e362bcbc2fe/nvidia_cudnn_cu12-9.25.1.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c71f4af9025e05e478cfb07f9c2a4c118346dac236e915f3e63769fbcb031cb3", size = 809606039, upload-time = "2026-08-28T23:39:06.797Z" },
@@ -1440,7 +1442,7 @@ name = "nvidia-cufft-cu12"
 version = "11.4.1.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" },
@@ -1452,9 +1454,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.7.5.82"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12" },
-    { name = "nvidia-cusparse-cu12" },
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" },
@@ -1466,7 +1468,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.5.10.65"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12" },
+    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" },
@@ -1496,7 +1498,7 @@ name = "nvidia-nvshmem-cu12"
 version = "3.7.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cuda-cccl-cu12" },
+    { name = "nvidia-cuda-cccl-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/88/a1/9ca9cc3338217311cee21f7ba795dad4f38b7a279a6da7d3d549258f60b1/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:98495de30994f1401a12cc8158ffacbeada44579f54185d75665381468056faa", size = 229963056, upload-time = "2026-07-17T19:23:44.067Z" },
@@ -1741,6 +1743,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" },
 ]
 
+[[package]]
+name = "pytest-timeout"
+version = "2.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"