jax-ml/jax · error · ValueError

trace_value requires a scalar value, got shape {value.shape}

Error message

trace_value requires a scalar value, got shape {value.shape}

What it means

trace_value in JAX Mosaic Pallas emits a debug trace of a value during kernel execution on TPU, but only scalar values can be traced. Any array with a non-empty shape is rejected at abstract-eval time.

Source

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

    pltpu.trace_value("my_x", x)
  """
  trace_value_p.bind(value, label=label)


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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Trace a scalar element, e.g. trace_value(x[0, 0]) or trace_value(x.flatten()[0])
2. Trace a reduction of the tile, e.g. trace_value(x.sum()) if a representative value suffices (must be i32/f32)

Example fix

# before
trace_value(block, label='block')
# after
trace_value(block[0, 0], label='block00')

assert not getattr(value, 'shape', None), 'trace_value needs a scalar'

def is_scalar_aval(value) -> bool:
    return not getattr(value, 'shape', None)

When it happens

Trigger: Calling trace_value(arr) on an array with ndim > 0, e.g. a Block of a tiled kernel or a vector produced inside the kernel body.

Common situations: Trying to trace a whole tile or vector for debugging instead of a single element; forgetting to index the block first (arr[0, 0]).

Related errors

• trace_value requires i32 or f32, got {value.dtype}
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

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