Issue 1:
Please:

Check for duplicate requests.

Describe your goal, and if possible provide a code snippet with a motivating example.

Issue 2:
Description

2023-05-29 05:39:19.492073: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:427] Loaded runtime CuDNN library: 8.5.0 but source was compiled with: 8.6.0. CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library. If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.
Traceback (most recent call last):
File "/t5x_home/t5x/train.py", line 834, in <module>
config_utils.run(main)
File "/t5x_home/t5x/config_utils.py", line 214, in run
gin_utils.run(main)
File "/t5x_home/t5x/gin_utils.py", line 129, in run
app.run(
File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/t5x_home/t5x/train.py", line 787, in main
_main(argv)
File "/t5x_home/t5x/train.py", line 829, in _main
train_using_gin()
File "/usr/local/lib/python3.8/dist-packages/gin/config.py", line 1605, in gin_wrapper
utils.augment_exception_message_and_reraise(e, err_str)
File "/usr/local/lib/python3.8/dist-packages/gin/utils.py", line 41, in augment_exception_message_and_reraise
raise proxy.with_traceback(exception.__traceback__) from None
File "/usr/local/lib/python3.8/dist-packages/gin/config.py", line 1582, in gin_wrapper
return fn(*new_args, **new_kwargs)
File "/t5x_home/t5x/train.py", line 265, in train
rng = random.PRNGKey(random_seed)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/random.py", line 137, in PRNGKey
key = prng.seed_with_impl(impl, seed)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 320, in seed_with_impl
return random_seed(seed, impl=impl)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 734, in random_seed
return random_seed_p.bind(seeds_arr, impl=impl)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/core.py", line 380, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/core.py", line 383, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/core.py", line 790, in process_primitive
return primitive.impl(*tracers, **params)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 746, in random_seed_impl
base_arr = random_seed_impl_base(seeds, impl=impl)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 751, in random_seed_impl_base
return seed(seeds)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 1419, in _rbg_seed
halfkey = threefry_seed(seed)
File "/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py", line 980, in threefry_seed
return _threefry_seed(seed)
jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.
In call to configurable 'train' (<function train at 0x7efc24d265e0>)

What jax/jaxlib version are you using?

jax 0.4.10 jaxlib-0.4.10+cuda11.cudnn86-cp38-cp38

Which accelerator(s) are you using?

GPU

Additional system info

ubuntu 22.04

cuda 11.8

cudnn 8.8

NVIDIA GPU info

No response