[jax_plugins]
cuda = jax_plugins.cuda
