How to solve wsl jax XlaRuntimeError

[NullXeronier]

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): WSL 2 Ubuntu 24.04.3 LTS noble
• Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib): flax 0.11.2 jax 0.7.2 jaxlib 0.7.2 / conda 25.7.0
• Python version: Python 3.12.4
• GPU/TPU model and memory: RTX 4090 VRAM 24GB / NVIDIA-SMI 580.95.02 / Driver Version: 581.42
• CUDA version (if applicable): 13.0

Problem you have encountered:

from flax import linen as nn Dense dummy random input value pass through nn Dense layer, nn Dense layer returns nothing and raise JaxRuntimeError, XlaRuntimeError: INTERNAL: an unsupported value or parameter was passed to the function

What you expected to happen:

flax linen nn Dense returns value like Tensorflow 2 keras layers Dense

Logs, error messages, etc:

---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[7], [line 9]
      7 x = jnp.ones((2, 4, 8), dtype=jnp.float32)
      8 model = TestDense(features=24)
----> [9] params = model.init(jax.random.PRNGKey(0), x)
     10 out = model.apply(params, x)
     11 print(out.shape)

    [... skipping hidden 9 frame]

Cell In[7], [line 5] 
      3 @nn.compact
      4 def __call__(self, x):
----> [5]     return nn.Dense(self.features)(x)

    [... skipping hidden 2 frame]

File ~/miniconda3/envs/kaggler/lib/python3.12/site-packages/flax/linen/linear.py:287, in Dense.__call__(self, inputs)
    285 else:
    286   dot_general = lax.dot_general
--> [287] y = dot_general(
    288   inputs,
    289   kernel,
    290   (((inputs.ndim - 1,), (0,)), ((), ())),
    291   precision=self.precision,
    292 )
    293 if bias is not None:
    294   y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))

    [... skipping hidden 18 frame]

File ~/miniconda3/envs/kaggler/lib/python3.12/site-packages/jax/_src/compiler.py:378, in backend_compile_and_load(backend, module, executable_devices, options, host_callbacks)
    369       return backend.compile_and_load(
    370           built_c,
    371           executable_devices=executable_devices,
    372           compile_options=options,
    373           host_callbacks=host_callbacks,
    374       )
    375     # Some backends don't have `host_callbacks` option yet
    376     # TODO(sharadmv): remove this fallback when all backends allow `compile`
    377     # to take in `host_callbacks`
--> [378] return backend.compile_and_load(
    379         built_c,
    380         executable_devices=executable_devices,
    381         compile_options=options,
    382     )
    383 except xc.XlaRuntimeError as e:
    384   for error_handler in _XLA_RUNTIME_ERROR_HANDLERS:

XlaRuntimeError: INTERNAL: an unsupported value or parameter was passed to the function

google/flax#5010

Steps to reproduce:

Whenever possible, please provide a minimal example. Please consider submitting it as a Colab link.
https://colab.research.google.com/drive/1Oo-g6pV5Sf7AjSVnd_RGteJXxb-xdfpz?usp=sharing

from flax import linen as nn
import jax.numpy as jnp
from chex import Array # chex 0.1.91
from jax import random
import jax

class TestDense(nn.Module):
    features: int
    @nn.compact
    def __call__(self, x):
        return nn.Dense(self.features)(x)

x = jnp.ones((2, 4, 8), dtype=jnp.float32)
model = TestDense(features=24)
params = model.init(jax.random.PRNGKey(0), x)
out = model.apply(params, x)
print(out.shape)

[NullXeronier]

solved seting xla arg
using #29031

unset XLA_FLAGS
export XLA_FLAGS="--xla_gpu_enable_cublaslt=true \
                  --xla_gpu_cublas_fallback=true \
                  --xla_gpu_enable_command_buffer=''" 
python - <<'PY'
import jax, jax.numpy as jnp
f = jax.jit(lambda x,y: x@y)
a=jnp.ones((256,256), jnp.float32); b=a
f(a,b).block_until_ready()
print("OK: cuBLASLt + fallback")
PY