What causes "Unsupported cast: {from_dtype} -> {to_dtype}" in jax-ml/jax?

Cast uint -> int32 first, then int32 -> float: x.astype(jnp.int32).astype(jnp.float32)

jax-ml/jax · error · NotImplementedError

Unsupported cast: {from_dtype} -> {to_dtype}

Error message

Unsupported cast: {from_dtype} -> {to_dtype}

What it means

The dtype-cast helper rejects casts it cannot lower: notably unsigned-to-float casts (uint -> float is unsupported on Mosaic TPU), after already handling unsigned->smaller-unsigned and unsigned widening. Any remaining combination falls through to this NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3149

      if from_bitwidth < 32:
        x = x.astype(jnp.float32)
    elif jnp.issubdtype(from_dtype, jnp.integer):
      if from_bitwidth < 32:
        x = x.astype(jnp.int32)
    return x != jnp.asarray(0, dtype=x.dtype)
  if jnp.issubdtype(from_dtype, jnp.signedinteger):
    if from_bitwidth < 32:
      x = x.astype(jnp.int32)
    if jnp.issubdtype(to_dtype, jnp.floating) and to_bitwidth < 32:
      x = x.astype(jnp.float32)
    return x.astype(to_dtype)
  if jnp.issubdtype(from_dtype, jnp.unsignedinteger):
    if from_bitwidth < 32:
      x = x.astype(jnp.uint32)
    # unsigned -> float is unsupported. We fall through and raise at the bottom.
    if not jnp.issubdtype(to_dtype, jnp.floating):
      return x.astype(to_dtype)
  raise NotImplementedError(f"Unsupported cast: {from_dtype} -> {to_dtype}")


@register_lowering_rule(
    lax.convert_element_type_p, kernel_types=[*tpu_core.CoreType]
)
def _convert_element_type_lowering_rule(
    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding
):
  del weak_type
  del sharding
  out_aval = ctx.avals_out[0]
  in_aval = ctx.avals_in[0]
  old_dtype = in_aval.dtype
  out_type = ctx.aval_to_ir_type(out_aval)

  if old_dtype == new_dtype:
    return x

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast uint -> int32 first, then int32 -> float: x.astype(jnp.int32).astype(jnp.float32)
Do the float conversion outside the kernel, pass float inputs in
Store source data as signed int8/int32

Example fix

// before
xf = x_u8.astype(jnp.float32)
// after
xf = x_u8.astype(jnp.int32).astype(jnp.float32)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax.numpy as jnp
def safe_cast(x, to):
    if jnp.issubdtype(x.dtype, jnp.unsignedinteger) and jnp.issubdtype(to, jnp.floating):
        return x.astype(jnp.int32).astype(to)
    return x.astype(to)

Type guard

def needs_uint_detour(from_dt, to_dt) -> bool:
    import jax.numpy as jnp
    return (jnp.issubdtype(from_dt, jnp.unsignedinteger)
            and jnp.issubdtype(to_dt, jnp.floating))

Prevention

Never cast uint directly to float in Mosaic kernels
Centralize casts in a helper
Convert data to float on the host before passing to kernels

When it happens

Trigger: convert_element_type from an unsigned integer dtype (e.g. uint8) to a floating dtype (float32/bfloat16) inside a Pallas Mosaic TPU kernel; also any other unhandled from->to pair.

Common situations: Preprocessing uint8 image data to float inside a Pallas kernel; normalizing quantized values with .astype(jnp.float32).

Related errors

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