Try rebasing on head as well, please. Sorry for the slow review.

Rebased again to resolve new conflicts.

There was a recent change that has an impact on this implementation too:

jax/jax_plugins/cuda/__init__.py

Lines 123 to 141 in 09d903f

	def _load_nvidia_libraries():
	"""Attempts to load NVIDIA's libraries.
	We prefer the Python packages, if present. If not, we fall back to loading
	them from LD_LIBRARY_PATH. By loading the libraries here, later lookups will
	find these copies."""
	_load("cuda_runtime", ["libcudart.so.12"])
	# cuda_nvrtc isn't directly a dependency of JAX, but CUDNN appears to need it
	# and at least in CUDA 12.9 has RUNPATHs misconfigured to refer to
	# nvidia/nvrtc instead of nvidia/cuda_nvrtc.
	_load("cuda_nvrtc", ["libnvrtc.so.12"])
	_load("cublas", ["libcublas.so.12", "libcublasLt.so.12"])
	_load("nccl", ["libnccl.so.2"])
	_load("cuda_cupti", ["libcupti.so.12"])
	_load("cusparse", ["libcusparse.so.12"])
	_load("cusolver", ["libcusolver.so.11"])
	_load("cufft", ["libcufft.so.11"])
	_load("nvshmem", ["libnvshmem_host.so.3"])
	_load("cudnn", ["libcudnn.so.9"])

I think you need something similar with try-except to load either CUDA12 or CUDA13 wheels.

Thanks. I think this one can be addressed in a follow-up if that's OK.

@olupton reran ci with latest changes, looks like there is still an issue in the cuda test and linter

@olupton reran ci with latest changes, looks like there is still an issue in the cuda test and linter

sorry, forgot to mention two more targets that need deps updated: jax_cuda_plugin_wheel_size_test and jax_cuda_pjrt_wheel_size_test.

Done, sorry about that. Tried to fix the mypy linter failure too.

Getting one additional linter error internally on jax.bzl. I'm going to try and patch it myself