{"record":{"id":"e1d88bd39b87edd1","repo":"jax-ml/jax","slug":"trace-value-requires-i32-or-f32-got-value-dtype","errorCode":null,"errorMessage":"trace_value requires i32 or f32, got {value.dtype}","messagePattern":"trace_value requires i32 or f32, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1182,"sourceCode":"\nclass TraceEffect(effects.Effect):\n  pass\n\n\ntrace_effect = TraceEffect()\neffects.control_flow_allowed_effects.add_type(TraceEffect)\npl_core.kernel_local_effects.add_type(TraceEffect)\n\n\n@trace_value_p.def_effectful_abstract_eval\ndef _trace_value_abstract_eval(value, *, label):\n  del label\n  if value.shape:\n    raise ValueError(\n        f\"trace_value requires a scalar value, got shape {value.shape}\"\n    )\n  if value.dtype not in (jnp.int32, jnp.float32):\n    raise ValueError(f\"trace_value requires i32 or f32, got {value.dtype}\")\n  return [], {trace_effect}\n\n\nclass MXUEffect(effects.Effect):\n  __str__ = lambda self: \"MXU\"\nmxu_effect = MXUEffect()\neffects.control_flow_allowed_effects.add_type(MXUEffect)\npl_core.kernel_local_effects.add_type(MXUEffect)\n\n\nmatmul_push_rhs_p = jax_core.Primitive(\"matmul_push_rhs\")\nmatmul_push_rhs_p.multiple_results = True\n\n\ndef matmul_push_rhs(\n    rhs: jax.Array,\n    staging_register: int,\n    mxu_index: int,","sourceCodeStart":1164,"sourceCodeEnd":1200,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1164-L1200","documentation":"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.","triggerScenarios":"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.","commonSituations":"Tracing a bfloat16 intermediate; tracing an int8/int64 accumulator or index value.","solutions":["Cast before tracing: trace_value(x.astype(jnp.float32))","For ints, cast to jnp.int32 first"],"exampleFix":"# before\ntrace_value(bf16_val, label='v')\n# after\ntrace_value(bf16_val.astype(jnp.float32), label='v')","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\nassert value.dtype in (jnp.int32, jnp.float32), f'cast {value.dtype} first'","typeGuard":"def traceable_dtype(x) -> bool:\n    return x.dtype in (jnp.int32, jnp.float32)","tryCatchPattern":null,"preventionTips":["Wrap tracing in a helper that casts to f32/i32 first"],"tags":["jax","pallas","tpu","debugging","dtype"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}