ITADN
jax-ml/jax/Issues

jax.nn.scaled_matmul raises NotImplementedError on CPU with no graceful error

#38813Openanna-researcher 创建于 2026-06-26
bug
A
anna-researchercommented
### Description `jax.nn.scaled_matmul` crashes immediately on CPU with an internal `NotImplementedError` rather than raising a clear user-facing error indicating the function requires GPU/cuDNN. ## Minimal reproducer ```python import jax import jax.numpy as jnp key = jax.random.PRNGKey(0) x = jax.random.normal(key, [3, 16, 16], dtype=jnp.bfloat16) out = jax.nn.scaled_matmul(x, x, x, x) ``` ## Error NotImplementedError: MLIR translation rule for primitive 'scaled_matmul' not found for platform cpu ## Expected behavior A clear error message stating that `scaled_matmul` requires a GPU with cuDNN support, not an internal MLIR lowering error. ### System info (python version, jaxlib version, accelerator, etc.) jax: 0.10.1 jaxlib: 0.10.1 numpy: 2.4.6 python: 3.12.13 | packaged by Anaconda, Inc. | (main, Mar 19 2026, 20:20:58) [GCC 14.3.0] device info: cpu-1, 1 local devices JAX_PLATFORMS=cpu
1 条评论