jax-ml/jax · error · NotImplementedError

cannot cast {src} to {dst_type}

Error message

cannot cast {src} to {dst_type}

What it means

This is the final fallback in _ir_cast: the cast between the given MLIR types has no lowering rule implemented in JAX's Triton Pallas backend. Only int<->int, float<->float (via F32 for f16/bf16), ptr<->ptr bitcasts and a few signed/unsigned combinations are handled; everything else falls through to this NotImplementedError.

Source

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

  if _is_triton_pointer_type(src_element_type) and isinstance(
      dst_element_type, ir.IntegerType
  ):
    if dst_element_type.width == 64:
      return tt_dialect.ptr_to_int(dst_type, src)
    elif dst_element_type.width == 1:
      x = _ir_cast(src, ir.IntegerType.get_signless(64), signed=signed)
      zero = _zeros_like(x)
      return _ir_cast(_not_equal(x, zero, signed=signed), dst_type, signed=signed)
  if isinstance(src_element_type, ir.IntegerType) and _is_triton_pointer_type(
      dst_element_type
  ):
    return tt_dialect.int_to_ptr(dst_type, src)
  if _is_triton_pointer_type(src_element_type) and _is_triton_pointer_type(
      dst_element_type
  ):
    return tt_dialect.bitcast(dst_type, src)

  raise NotImplementedError(f"cannot cast {src} to {dst_type}")


@register_lowering(lax.convert_element_type_p)
def _convert_element_type_lowering_rule(
    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding
):
  [x_aval] = ctx.avals_in
  x = _ensure_ir_value(x, x_aval)
  if new_dtype == x_aval.dtype:
    return x
  cc = ctx.context.compute_capability
  return _cast(x, x_aval.dtype, new_dtype, compute_capability=cc)


@register_lowering(lax.select_n_p)
def select_n_lowering_rule(ctx: LoweringRuleContext, pred, x, y):
  pred_aval, a_aval, b_aval = ctx.avals_in
  [out_aval] = ctx.avals_out

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the cast into supported steps (e.g. via f32 for float types, or explicit signed casts)
2. Avoid pointer<->numeric casts in the Python kernel body; compute integer offsets instead
3. Print src/dst MLIR types at the failure site to identify which branch is missing and restructure the kernel accordingly
4. File/check a JAX GitHub issue — missing cast combinations are implemented incrementally

Example fix

// before
p = x.astype some_ptr_type  # float -> pointer

// after
off = x.astype(jnp.int32)
p = base_ptr + off  # use pointer arithmetic on integer offsets

try:
    lowered = kernel_compile_and_run(...)
except NotImplementedError as e:
    if 'cannot cast' in str(e):
        # inspect e for src/dst types and split the cast via f32
        x = x.astype(jnp.float32).astype(target)

When it happens

Trigger: Casting between type classes the lowering does not handle, e.g. float to pointer, integer to float (in some paths), complex types, or casts involving Triton pointer types not matching the int_to_ptr/ptr bitcast branches. Typically surfaces from convert_element_type, load/store `other` values, or index computation.

Common situations: Exotic casts inside Pallas kernels (e.g. reinterpret-style tricks that worked in Triton-lang), complex64 kernels, bool/int1 handling edge cases, or version mismatches where the kernel was written against newer lowering capabilities.

Related errors

• Unsupported aval type: {aval}, {type(aval)}, {t}
• cannot cast from `{dtype_name}`
• cannot cast to `{dtype_name}`
• Only 2-argument concatenate is supported.
• Only concatenate along the last dimension is supported.

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