MAINT: bump JAX to 0.7.2

Author: lucascolley

pixi.lock

@@ -153,15 +153,13 @@ dask-core = ">=2025.12.0" # No distributed, tornado, etc.
 sparse = ">=0.17.0"
 
 [tool.pixi.feature.backends.target.linux-64.dependencies]
-# On CPU use >=0.7.0
-# On GPU, use 0.6.0 (0.6.2 and 0.7.0 both segfault); see jaxlib pin below.
-jax = ">=0.6.0"
+jax = ">=0.7.2"
 
 [tool.pixi.feature.backends.target.osx-64.dependencies]
-jax = ">=0.6.0"
+jax = ">=0.7.2"
 
 [tool.pixi.feature.backends.target.osx-arm64.dependencies]
-jax = ">=0.6.0"
+jax = ">=0.7.2"
 
 [tool.pixi.feature.backends.target.win-64.dependencies]
 # jax = "*"  # unavailable
@@ -175,23 +173,17 @@ jax = ">=0.6.0"
 [tool.pixi.feature.cuda-backends]
 system-requirements = { cuda = "12" }
 
-[tool.pixi.feature.cuda-backends.target.linux-64.dependencies]
+[tool.pixi.feature.cuda-backends.target.linux.dependencies]
 cupy = ">=13.6.0"
-# JAX 0.6.2 and 0.7.0 segfault on CUDA
-jaxlib = { version = ">=0.6.0,!=0.6.2,!=0.7.0", build = "cuda12*" }
+jaxlib = { version = ">=0.7.2", build = "cuda12*" }
 pytorch = { version = ">=2.9.1", build = "cuda12*" }
 
-[tool.pixi.feature.cuda-backends.target.osx-64.dependencies]
+[tool.pixi.feature.cuda-backends.target.osx.dependencies]
 # cupy = "*"  # unavailable
 # jaxlib = { version = "*", build = "cuda12*" }  # unavailable
 # pytorch = { version = "*", build = "cuda12*" }  # unavailable
 
-[tool.pixi.feature.cuda-backends.target.osx-arm64.dependencies]
-# cupy = "*"  # unavailable
-# jaxlib = { version = "*", build = "cuda12*" }  # unavailable
-# pytorch = { version = "*", build = "cuda12*" }  # unavailable
-
-[tool.pixi.feature.cuda-backends.target.win-64.dependencies]
+[tool.pixi.feature.cuda-backends.target.win.dependencies]
 cupy = ">=13.6.0"
 # jaxlib = { version = "*", build = "cuda12*" }  # unavailable
 pytorch = { version = ">=2.9.1", build = "cuda12*" }