diff --git a/pixi.lock b/pixi.lock
index 3ce120cf..76ae4833 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -142,8 +142,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/ipython-8.36.0-pyh907856f_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/isort-6.0.1-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.5.2-cpu_py310hc96afab_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.6.0-cpu_py310hc96afab_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/linux-64/ld_impl_linux-64-2.43-h712a8e2_4.conda
@@ -315,8 +315,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/ipython-8.36.0-pyh907856f_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/isort-6.0.1-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.5.2-cpu_py310h22b337c_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.6.0-cpu_py310h22b337c_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libabseil-20250127.1-cxx17_h0e468a2_0.conda
@@ -481,8 +481,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/ipython-8.36.0-pyh907856f_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/isort-6.0.1-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.5.2-cpu_py310h2c532f2_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.6.0-cpu_py310h2c532f2_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libabseil-20250127.1-cxx17_h07bc746_0.conda
@@ -1022,8 +1022,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/ipython-8.36.0-pyh907856f_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/isort-6.0.1-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.5.2-cpu_py310h22b337c_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.6.0-cpu_py310h22b337c_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libabseil-20250127.1-cxx17_h0e468a2_0.conda
@@ -1188,8 +1188,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/ipython-8.36.0-pyh907856f_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/isort-6.0.1-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.5.2-cpu_py310h2c532f2_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.6.0-cpu_py310h2c532f2_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libabseil-20250127.1-cxx17_h07bc746_0.conda
@@ -2508,8 +2508,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/linux-64/icu-75.1-he02047a_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.5.2-cpu_py310hc96afab_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.6.0-cpu_py310hc96afab_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/linux-64/ld_impl_linux-64-2.43-h712a8e2_4.conda
       - conda: https://prefix.dev/conda-forge/linux-64/libabseil-20250127.1-cxx17_hbbce691_0.conda
@@ -2606,8 +2606,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/osx-64/icu-75.1-h120a0e1_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.5.2-cpu_py310h22b337c_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.6.0-cpu_py310h22b337c_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libabseil-20250127.1-cxx17_h0e468a2_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libblas-3.9.0-20_osx64_mkl.conda
@@ -2696,8 +2696,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/hypothesis-6.131.17-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.5.2-cpu_py310h2c532f2_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.6.0-cpu_py310h2c532f2_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libabseil-20250127.1-cxx17_h07bc746_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libblas-3.9.0-31_h10e41b3_openblas.conda
@@ -3014,8 +3014,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/osx-64/icu-75.1-h120a0e1_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.5.2-cpu_py310h22b337c_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.6.0-cpu_py310h22b337c_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libabseil-20250127.1-cxx17_h0e468a2_0.conda
       - conda: https://prefix.dev/conda-forge/osx-64/libblas-3.9.0-20_osx64_mkl.conda
@@ -3104,8 +3104,8 @@ environments:
       - conda: https://prefix.dev/conda-forge/noarch/hypothesis-6.131.17-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.6.1-pyha770c72_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.0.0-pyhd8ed1ab_1.conda
-      - conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.5.2-cpu_py310h2c532f2_2.conda
+      - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
+      - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.6.0-cpu_py310h2c532f2_0.conda
       - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libabseil-20250127.1-cxx17_h07bc746_0.conda
       - conda: https://prefix.dev/conda-forge/osx-arm64/libblas-3.9.0-31_h10e41b3_openblas.conda
@@ -3915,7 +3915,7 @@ packages:
 - pypi: ./
   name: array-api-extra
   version: 0.8.0.dev0
-  sha256: a3296bd31e0640360736388accc9305ee606fe6cafb22b8fdd7492fa96a57cf8
+  sha256: 372435cdb1289323015e8675e2d91b881fec3f5f072e3ccedff93528b4800f2c
   requires_dist:
   - array-api-compat>=1.12.0,<2
   requires_python: '>=3.10'
@@ -5527,25 +5527,6 @@ packages:
   - pkg:pypi/isort?source=hash-mapping
   size: 75118
   timestamp: 1746190379093
-- conda: https://prefix.dev/conda-forge/noarch/jax-0.5.2-pyhd8ed1ab_0.conda
-  sha256: be7644c955cd4be330a13a8f64c0b73d520f8b3ab6bb64b8b1d3a17945345684
-  md5: f19f3d281603af8e67d533dbeac279ce
-  depends:
-  - importlib-metadata >=4.6
-  - jaxlib >=0.5.1,<=0.5.2
-  - ml_dtypes >=0.4.0
-  - numpy >=1.25
-  - opt_einsum
-  - python >=3.10
-  - scipy >=1.11.1
-  constrains:
-  - cudnn >=9.2.1.18,<10.0
-  license: Apache-2.0
-  license_family: APACHE
-  purls:
-  - pkg:pypi/jax?source=hash-mapping
-  size: 1556886
-  timestamp: 1741182198677
 - conda: https://prefix.dev/conda-forge/noarch/jax-0.6.0-pyhd8ed1ab_0.conda
   sha256: 573a5582dfba84a8f67c351b6218cb9579cb8d0f6d4b4186a806852111d4a6f1
   md5: bd364feb12c744cf5c60e1e5b586171b
@@ -5565,9 +5546,9 @@ packages:
   - pkg:pypi/jax?source=hash-mapping
   size: 1538293
   timestamp: 1748688029463
-- conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.5.2-cpu_py310hc96afab_2.conda
-  sha256: bc9fc4a0bc9b43b0799c3320b62bd5e486e57ac70396ee5236a70a87018723b6
-  md5: 2bbb84513ad8ff868dd21ef5e5661aaa
+- conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.6.0-cpu_py310hc96afab_0.conda
+  sha256: c54f8e339df547924c468e7632981f553cd764f75002e171900d9df187898e4b
+  md5: e4e56540b3315f1ac79440a1465f5182
   depends:
   - __glibc >=2.17,<3.0.a0
   - libabseil * cxx17*
@@ -5583,13 +5564,13 @@ packages:
   - python_abi 3.10.* *_cp310
   - scipy >=1.9
   constrains:
-  - jax >=0.5.2
+  - jax >=0.6.0
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/jaxlib?source=hash-mapping
-  size: 69060155
-  timestamp: 1747094262205
+  size: 60603285
+  timestamp: 1748663385056
 - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.6.0-cuda126py310hec873cc_200.conda
   sha256: 5a52619f5e6d40d5c8e8223a5ec113d5bb097456656bc5668ab24df6b75f69ae
   md5: 1b39986ae9b1bcfddc720ef10bb67420
@@ -5635,9 +5616,9 @@ packages:
   - pkg:pypi/jaxlib?source=hash-mapping
   size: 146820753
   timestamp: 1748663708635
-- conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.5.2-cpu_py310h22b337c_2.conda
-  sha256: bb292699e4523002e868c9b56379574c74ef82f718ef574d27fefa0a882ae0e0
-  md5: 2f10f0eb0ddf22befa0d8c750d2df551
+- conda: https://prefix.dev/conda-forge/osx-64/jaxlib-0.6.0-cpu_py310h22b337c_0.conda
+  sha256: 1ea4b15f45eeecd939270d93675f7b7361c9a6cd7a830d361920855a10573b59
+  md5: 25c084b06a0c200d1cc095b59f24a757
   depends:
   - __osx >=10.15
   - libabseil * cxx17*
@@ -5652,16 +5633,16 @@ packages:
   - python_abi 3.10.* *_cp310
   - scipy >=1.9
   constrains:
-  - jax >=0.5.2
+  - jax >=0.6.0
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/jaxlib?source=hash-mapping
-  size: 66014116
-  timestamp: 1747095379031
-- conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.5.2-cpu_py310h2c532f2_2.conda
-  sha256: f4d9d8e57db13d09eb3f83f3c57bceef327689075b434dff44137045dd43da33
-  md5: e9f65c12435b42a8094701861d02f25b
+  size: 62834580
+  timestamp: 1748654408664
+- conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.6.0-cpu_py310h2c532f2_0.conda
+  sha256: a35d81d3ac366fcff2d5cba19ae5452b9aa8fd87ea4ace34a737f2fa88a85826
+  md5: 1ab7fe4ca08fb32518bdeea813ef6bef
   depends:
   - __osx >=11.0
   - libabseil * cxx17*
@@ -5677,13 +5658,13 @@ packages:
   - python_abi 3.10.* *_cp310
   - scipy >=1.9
   constrains:
-  - jax >=0.5.2
+  - jax >=0.6.0
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/jaxlib?source=hash-mapping
-  size: 54831498
-  timestamp: 1747096397162
+  size: 51702562
+  timestamp: 1748652584379
 - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda
   sha256: 92c4d217e2dc68983f724aa983cca5464dcb929c566627b26a2511159667dba8
   md5: a4f4c5dc9b80bc50e0d3dc4e6e8f1bd9
diff --git a/pyproject.toml b/pyproject.toml
index b710e5d2..84dc07b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -139,13 +139,13 @@ dask-core = ">=2025.5.0"  # No distributed, tornado, etc.
 sparse = ">=0.16.0"
 
 [tool.pixi.feature.backends.target.linux-64.dependencies]
-jax = ">=0.5.2"
+jax = ">=0.5.3"
 
 [tool.pixi.feature.backends.target.osx-64.dependencies]
-jax = ">=0.5.2"
+jax = ">=0.5.3"
 
 [tool.pixi.feature.backends.target.osx-arm64.dependencies]
-jax = ">=0.5.2"
+jax = ">=0.5.3"
 
 [tool.pixi.feature.backends.target.win-64.dependencies]
 # jax = "*"  # unavailable