jax-ml/jax · error · ValueError

trace_value requires i32 or f32, got {value.dtype}

Error message

trace_value requires i32 or f32, got {value.dtype}

What it means

trace_value in JAX Mosaic Pallas only supports int32 and float32 values because the TPU tracing hardware only handles those widths. Other dtypes (bf16, f64, int8, etc.) are rejected at abstract-eval time.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1182

class TraceEffect(effects.Effect):
  pass


trace_effect = TraceEffect()
effects.control_flow_allowed_effects.add_type(TraceEffect)
pl_core.kernel_local_effects.add_type(TraceEffect)


@trace_value_p.def_effectful_abstract_eval
def _trace_value_abstract_eval(value, *, label):
  del label
  if value.shape:
    raise ValueError(
        f"trace_value requires a scalar value, got shape {value.shape}"
    )
  if value.dtype not in (jnp.int32, jnp.float32):
    raise ValueError(f"trace_value requires i32 or f32, got {value.dtype}")
  return [], {trace_effect}


class MXUEffect(effects.Effect):
  __str__ = lambda self: "MXU"
mxu_effect = MXUEffect()
effects.control_flow_allowed_effects.add_type(MXUEffect)
pl_core.kernel_local_effects.add_type(MXUEffect)


matmul_push_rhs_p = jax_core.Primitive("matmul_push_rhs")
matmul_push_rhs_p.multiple_results = True


def matmul_push_rhs(
    rhs: jax.Array,
    staging_register: int,
    mxu_index: int,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast before tracing: trace_value(x.astype(jnp.float32))
2. For ints, cast to jnp.int32 first

Example fix

# before
trace_value(bf16_val, label='v')
# after
trace_value(bf16_val.astype(jnp.float32), label='v')

import jax.numpy as jnp
assert value.dtype in (jnp.int32, jnp.float32), f'cast {value.dtype} first'

def traceable_dtype(x) -> bool:
    return x.dtype in (jnp.int32, jnp.float32)

When it happens

Trigger: Calling trace_value(x) where x.dtype is not jnp.int32 or jnp.float32 — common with bfloat16 blocks which are the default in TPU kernels.

Common situations: Tracing a bfloat16 intermediate; tracing an int8/int64 accumulator or index value.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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