XLA error when running rng = jax.random.PRNGKey(seed = 42)

[PabloAMC]

I am getting a weird XLA error when using rng = jax.random.PRNGKey(seed = 42) in a jax program in google colab

!pip install --upgrade pip
!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install flax
!pip install optax

from google.colab import drive
drive.mount('/content/gdrive')

%cd "gdrive/MyDrive/project/"
import sys
sys.path.append('gdrive/MyDrive/project/')

%run project/main.py training

The error I get is

[/content/gdrive/MyDrive/project/functional.py](https://localhost:8080/#) in generate_weights(self, folder, num_layers_with_dm_parameters)
    302             return params
    303 
--> 304         rng = jax.random.PRNGKey(seed = 42)
    305         example_input = jax.random.normal(rng, shape=(1, 11))
    306         params = self.init(rng, example_input)

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr

The full trace is

[/content/gdrive/MyDrive/project/functional.py](https://localhost:8080/#) in generate_weights(self, folder, num_layers_with_dm_parameters)
    302             return params
    303 
--> 304         rng = jax.random.PRNGKey(seed = 42)
    305         example_input = jax.random.normal(rng, shape=(1, 11))
    306         params = self.init(rng, example_input)

[/usr/local/lib/python3.8/dist-packages/jax/_src/random.py](https://localhost:8080/#) in PRNGKey(seed)
    134     raise TypeError("PRNGKey accepts a scalar seed, but was given an array of"
    135                     f"shape {np.shape(seed)} != (). Use jax.vmap for batching")
--> 136   key = prng.seed_with_impl(impl, seed)
    137   return _return_prng_keys(True, key)
    138 

[/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py](https://localhost:8080/#) in seed_with_impl(impl, seed)
    265 # TODO(frostig): remove, rerouting callers directly to random_seed
    266 def seed_with_impl(impl: PRNGImpl, seed: int) -> PRNGKeyArray:
--> 267   return random_seed(seed, impl=impl)
    268 
    269 

[/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py](https://localhost:8080/#) in random_seed(seeds, impl)
    568   else:
    569     seeds_arr = jnp.asarray(seeds)
--> 570   return random_seed_p.bind(seeds_arr, impl=impl)
    571 
    572 random_seed_p = core.Primitive('random_seed')

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in bind(self, *args, **params)
    341     assert (not config.jax_enable_checks or
    342             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--> 343     return self.bind_with_trace(find_top_trace(args), args, params)
    344 
    345   def bind_with_trace(self, trace, args, params):

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in bind_with_trace(self, trace, args, params)
    344 
    345   def bind_with_trace(self, trace, args, params):
--> 346     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    347     return map(full_lower, out) if self.multiple_results else full_lower(out)
    348 

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in process_primitive(self, primitive, tracers, params)
    726 
    727   def process_primitive(self, primitive, tracers, params):
--> 728     return primitive.impl(*tracers, **params)
    729 
    730   def process_call(self, primitive, f, tracers, params):

[/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py](https://localhost:8080/#) in random_seed_impl(seeds, impl)
    580 @random_seed_p.def_impl
    581 def random_seed_impl(seeds, *, impl):
--> 582   base_arr = random_seed_impl_base(seeds, impl=impl)
    583   return PRNGKeyArray(impl, base_arr)
    584 

[/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py](https://localhost:8080/#) in random_seed_impl_base(seeds, impl)
    585 def random_seed_impl_base(seeds, *, impl):
    586   seed = iterated_vmap_unary(seeds.ndim, impl.seed)
--> 587   return seed(seeds)
    588 
    589 def random_seed_lowering(ctx, seeds, *, impl):

[/usr/local/lib/python3.8/dist-packages/jax/_src/prng.py](https://localhost:8080/#) in threefry_seed(seed)
    820   convert = lambda k: lax.reshape(lax.convert_element_type(k, np.uint32), [1])
    821   k1 = convert(
--> 822       lax.shift_right_logical(seed, lax_internal._const(seed, 32)))
    823   with jax.numpy_dtype_promotion('standard'):
    824     # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit

[/usr/local/lib/python3.8/dist-packages/jax/_src/lax/lax.py](https://localhost:8080/#) in shift_right_logical(x, y)
    509 def shift_right_logical(x: ArrayLike, y: ArrayLike) -> Array:
    510   r"""Elementwise logical right shift: :math:`x \gg y`."""
--> 511   return shift_right_logical_p.bind(x, y)
    512 
    513 def eq(x: ArrayLike, y: ArrayLike) -> Array:

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in bind(self, *args, **params)
    341     assert (not config.jax_enable_checks or
    342             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--> 343     return self.bind_with_trace(find_top_trace(args), args, params)
    344 
    345   def bind_with_trace(self, trace, args, params):

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in bind_with_trace(self, trace, args, params)
    344 
    345   def bind_with_trace(self, trace, args, params):
--> 346     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    347     return map(full_lower, out) if self.multiple_results else full_lower(out)
    348 

[/usr/local/lib/python3.8/dist-packages/jax/_src/core.py](https://localhost:8080/#) in process_primitive(self, primitive, tracers, params)
    726 
    727   def process_primitive(self, primitive, tracers, params):
--> 728     return primitive.impl(*tracers, **params)
    729 
    730   def process_call(self, primitive, f, tracers, params):

[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in apply_primitive(prim, *args, **params)
    120 def apply_primitive(prim, *args, **params):
    121   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 122   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
    123                                         **params)
    124   return compiled_fun(*args)

[/usr/local/lib/python3.8/dist-packages/jax/_src/util.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    251         return f(*args, **kwargs)
    252       else:
--> 253         return cached(config._trace_context(), *args, **kwargs)
    254 
    255     wrapper.cache_clear = cached.cache_clear

[/usr/local/lib/python3.8/dist-packages/jax/_src/util.py](https://localhost:8080/#) in cached(_, *args, **kwargs)
    244     @functools.lru_cache(max_size)
    245     def cached(_, *args, **kwargs):
--> 246       return f(*args, **kwargs)
    247 
    248     @functools.wraps(f)

[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in xla_primitive_callable(prim, *arg_specs, **params)
    199     else:
    200       return out,
--> 201   compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
    202                                     prim.name, donated_invars, False, *arg_specs)
    203   if not prim.multiple_results:

[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in _xla_callable_uncached(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)
    352                                    False, keep_unused, *arg_specs)
    353     allow_prop = [True] * len(computation.compile_args['global_out_avals'])
--> 354     return computation.compile(_allow_propagation_to_outputs=allow_prop).unsafe_call
    355   else:
    356     return lower_xla_callable(fun, device, backend, name, donated_invars, False,

[/usr/local/lib/python3.8/dist-packages/jax/_src/interpreters/pxla.py](https://localhost:8080/#) in compile(self, _allow_propagation_to_outputs, _allow_compile_replicated)
   3201               _allow_compile_replicated: bool = True) -> MeshExecutable:
   3202     if self._executable is None:
-> 3203       executable = self._compile_unloaded(
   3204           _allow_propagation_to_outputs, _allow_compile_replicated)
   3205       if isinstance(executable, UnloadedMeshExecutable):

[/usr/local/lib/python3.8/dist-packages/jax/_src/interpreters/pxla.py](https://localhost:8080/#) in _compile_unloaded(self, _allow_propagation_to_outputs, _allow_compile_replicated)
   3172       return MeshExecutable.from_trivial_jaxpr(**self.compile_args)
   3173     else:
-> 3174       return UnloadedMeshExecutable.from_hlo(
   3175           self._name,
   3176           self._hlo,

[/usr/local/lib/python3.8/dist-packages/jax/_src/interpreters/pxla.py](https://localhost:8080/#) in from_hlo(name, computation, mesh, global_in_avals, global_out_avals, in_shardings, out_shardings, spmd_lowering, tuple_args, in_is_global, auto_spmd_lowering, _allow_propagation_to_outputs, _allow_compile_replicated, unordered_effects, ordered_effects, host_callbacks, keepalive, kept_var_idx, backend, device_assignment, committed, pmap_nreps)
   3454                                      "in {elapsed_time} sec",
   3455                                      event=dispatch.BACKEND_COMPILE_EVENT):
-> 3456         xla_executable = dispatch.compile_or_get_cached(
   3457             backend, computation, compile_options, host_callbacks)
   3458 

[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in compile_or_get_cached(backend, computation, compile_options, host_callbacks)
   1079       return compiled
   1080 
-> 1081   return backend_compile(backend, serialized_computation, compile_options,
   1082                          host_callbacks)
   1083 

[/usr/local/lib/python3.8/dist-packages/jax/_src/profiler.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    312   def wrapper(*args, **kwargs):
    313     with TraceAnnotation(name, **decorator_kwargs):
--> 314       return func(*args, **kwargs)
    315     return wrapper
    316   return wrapper

[/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in backend_compile(backend, built_c, options, host_callbacks)
   1024   # TODO(sharadmv): remove this fallback when all backends allow `compile`
   1025   # to take in `host_callbacks`
-> 1026   return backend.compile(built_c, compile_options=options)
   1027 
   1028 _ir_dump_counter = itertools.count()

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:626) dnn != nullptr

[PabloAMC]

Solved in #14464 XLA Error when calling rng = jax.random.PRNGKey(0) in google colab
using

!pip install --upgrade "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html