What causes "cannot cast {src} tp {dst_type}" in jax-ml/jax?

First upcast the float8 value to float32, then cast to int: x.astype(jnp.float32).astype(jnp.int32)

jax-ml/jax · error · NotImplementedError

cannot cast {src} tp {dst_type}

Error message

cannot cast {src} tp {dst_type}

What it means

_float_int_cast lowers float→int casts, but only accepts source floats of bf16/f16/f32/f64. Any other source element type (float8 variants, custom float types) fails with 'cannot cast {src} tp {dst_type}' (note the 'tp' typo in the message). This path also implements bool casts via a not-equal comparison against zero.

Source

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

  if dst_element_type.width == 1:
    return _not_equal(src, _zeros_like(src), signed=signed)

  if src_element_type.width == dst_element_type.width:
    return arith_dialect.bitcast(dst_type, src)
  elif src_element_type.width > dst_element_type.width:
    return arith_dialect.trunci(dst_type, src)
  elif signed and src_element_type.width != 1:
    return arith_dialect.extsi(dst_type, src)
  else:
    return arith_dialect.extui(dst_type, src)


def _float_int_cast(
    src: ir.Value, dst_type: ir.Type, *, signed: bool
) -> ir.Value:
  src_element_type = _element_type(src.type)
  if not isinstance(src_element_type, (ir.BF16Type, ir.F16Type, ir.F32Type, ir.F64Type)):
    raise NotImplementedError(f"cannot cast {src} tp {dst_type}")
  dst_element_type = ir.IntegerType(_element_type(dst_type))
  if dst_element_type.width == 1:
    return _not_equal(src, _zeros_like(src), signed=signed)
  else:
    # We clamp the float value to the min/max integer destination value
    # in order to match JAX/XLA casting behavior. Note that this differs
    # from numpy casting behavior.
    if signed:
      maxint = 2**(dst_element_type.width-1) - 1
      minint = -2**(dst_element_type.width-1)
    else:
      maxint = 2**dst_element_type.width - 1
      minint = 0
    src = arith_dialect.minimumf(src, _full(src.type, maxint))
    src = arith_dialect.maximumf(src, _full(src.type, minint))
    if signed:
      return arith_dialect.fptosi(dst_type, src)
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

First upcast the float8 value to float32, then cast to int: x.astype(jnp.float32).astype(jnp.int32)
Restructure the kernel to keep values in supported dtypes (f32 accumulators) and convert at the boundaries
Upgrade JAX / check for float8 cast support in pallas triton lowering and file an issue with the exact dtype pair

Example fix

// before
i = x_f8.astype(jnp.int32)
// after
i = x_f8.astype(jnp.float32).astype(jnp.int32)

Defensive patterns

Strategy: fallback

Validate before calling

TRITON_FLOATS = (jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64)
def cast_int_safe(x, dt):
    if x.dtype not in TRITON_FLOATS:
        x = x.astype(jnp.float32)
    return x.astype(dt)

Type guard

def is_triton_castable_float(dt):
    return dt in (jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64)

Try / catch

try:
    out = pallas_kernel(x)
except NotImplementedError as e:
    if 'cannot cast' in str(e):
        x = x.astype(jnp.float32)
        out = pallas_kernel(x)
    else:
        raise

Prevention

Upcast float8 values to f32 before any int/bool cast in Triton kernels
Keep quantize/dequantize steps in f32 and convert at kernel boundaries

When it happens

Trigger: Casting a float8 tensor (e.g. e4m3/e5m2 from ml_dtypes) to an integer or bool inside a Triton Pallas kernel: x.astype(jnp.int32) or bool(x_f8); any custom FloatType outside the four supported ones reaching _ir_cast.

Common situations: Mixed-precision float8 kernels that quantize to int; dequantize/quantize pipelines written for XLA being ported to Pallas/Triton; ml_dtypes float8 usage becoming more common on H100/Blackwell GPUs.

Related errors

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