jax-ml/jax · error · NotImplementedError

cannot cast from `{dtype_name}`

Error message

cannot cast from `{dtype_name}`

What it means

Raised by the Triton Pallas lowering path in JAX when a cast (convert_element_type) has a source dtype that is not supported by the GPU's compute capability. The _UNSUPPORTED_CAST_DTYPES table lists dtypes (e.g. float8 variants, bfloat16 on older GPUs) whose is_supported(compute_capability) check fails, so the kernel cannot be lowered to Triton.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1685

  if isinstance(src.type, ir.RankedTensorType) and not isinstance(
      dst_type, ir.RankedTensorType
  ):
    src_type = ir.RankedTensorType(src.type)
    dst_type = ir.RankedTensorType.get(
        src_type.shape,
        dst_type,
        src_type.encoding,
    )
  if src.type == dst_type:
    return src

  src_element_type = _element_type(src.type)
  dst_element_type = _element_type(dst_type)

  for dtype, dtype_name, is_supported in _UNSUPPORTED_CAST_DTYPES:
    if isinstance(src_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast from `{dtype_name}`")
    if isinstance(dst_element_type, dtype):
      if not is_supported(compute_capability):
        raise NotImplementedError(f"cannot cast to `{dtype_name}`")

  if isinstance(src_element_type, (ir.F16Type, ir.BF16Type)) and not isinstance(
      dst_element_type, ir.F32Type
  ):
    return _ir_cast(
        _ir_cast(src, ir.F32Type.get(), signed=False),
        dst_type, signed=False, dst_signed=dst_signed
    )

  if isinstance(src_element_type, ir.FloatType) and isinstance(
      dst_element_type, ir.FloatType
  ):
    return _float_float_cast(src, dst_type)

  if isinstance(src_element_type, ir.IntegerType) and isinstance(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check your GPU compute capability and avoid fp8/bf16 source dtypes on unsupported hardware (fp8 needs Hopper/Ada+, sm89/sm90)
2. Convert the tensor to f32/f16 before passing it into the Triton Pallas kernel
3. Update JAX — support for more dtypes/casts is added over time in the Triton lowering
4. If the dtype is essential, fall back to the standard XLA backend instead of the Triton Pallas backend

Example fix

// before
out = lax.convert_element_type(x, jnp.float8_e4m3fn)  # on sm80

// after
out = lax.convert_element_type(lax.convert_element_type(x, jnp.float16), jnp.float8_e4m3fn)  # or run on sm89+

from jax._src.pallas.triton import helpers
# Before the kernel, gate on compute capability for exotic dtypes
import jax
cc = jax.devices()[0].compute_capability
FP8_OK = float(cc) >= 8.9
assert FP8_OK or x.dtype not in (jnp.float8_e4m3fn, jnp.float8_e5m2), 'fp8 unsupported on this GPU'

def cast_supported_on_device(dtype, cc: float) -> bool:
    if dtype in (jnp.float8_e4m3fn, jnp.float8_e5m2):
        return cc >= 8.9
    if dtype == jnp.bfloat16:
        return cc >= 8.0
    return True

try:
    kernel_lowered = triton_kernel(...)  
except NotImplementedError as e:
    if 'cannot cast' in str(e):
        x = x.astype(jnp.float16)  # widen and retry

When it happens

Trigger: Calling lax.convert_element_type (or a Pallas kernel that loads/stores/casts) with a source dtype such as f8e4m3/f8e5m2/bfloat16 on a GPU whose compute capability (e.g. pre-Hopper/Ada for fp8) does not support that type; also reached via _load/_store/_compute_offsets_from_indices which internally cast indices.

Common situations: Running fp8 kernels on Ampere or older GPUs; bf16 on very old GPUs; assuming JAX's default XLA path supports a dtype that the experimental Mosaic/Triton Pallas backend does not.

Related errors

• cannot cast to `{dtype_name}`
• Unsigned integer dtype {aval.dtype} is not supported for dot
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• cannot cast {src} to {dst_type}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/442ffdbfc7b7c66d. Report an issue: GitHub.