Revert JAX and make torch explicit too

Author: crusaderky

pixi.lock

@@ -154,23 +154,34 @@ jax = "*"
 [tool.pixi.feature.backends.target.win-64.dependencies]
 # jax = "*"  # unavailable
 
-# Backends that require a GPU host and a CUDA driver
-# Note: JAX and PyTorch automatically install CUDA variants
-# thanks to the `system-requirements` below.
+# Backends that require a GPU host and a CUDA driver.
+# Note that JAX and PyTorch automatically prefer CUDA variants
+# thanks to the `system-requirements` below, *if available*.
+# We request them explicitly below to ensure that we don't
+# quietly revert to CPU-only in the future, e.g. when CUDA 13
+# is released and CUDA 12 builds are dropped upstream.
 [tool.pixi.feature.cuda-backends]
 system-requirements = { cuda = "12" }
 
 [tool.pixi.feature.cuda-backends.target.linux-64.dependencies]
 cupy = "*"
+jaxlib = { version = "*", build = "cuda12*" }
+pytorch = { version = "*", build = "cuda12*" }
 
 [tool.pixi.feature.cuda-backends.target.osx-64.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]
 cupy = "*"
+# jaxlib = { version = "*", build = "cuda12*" }  # unavailable
+pytorch = { version = "*", build = "cuda12*" }
 
 [tool.pixi.environments]
 default = { features = ["py313"], solve-group = "py313" }