{"record":{"id":"3ebf3c8845849e6d","repo":"jax-ml/jax","slug":"trace-value-requires-a-scalar-value-got-shape-va","errorCode":null,"errorMessage":"trace_value requires a scalar value, got shape {value.shape}","messagePattern":"trace_value requires a scalar value, got shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1178,"sourceCode":"    pltpu.trace_value(\"my_x\", x)\n  \"\"\"\n  trace_value_p.bind(value, label=label)\n\n\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","sourceCodeStart":1160,"sourceCodeEnd":1196,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1160-L1196","documentation":"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.","triggerScenarios":"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.","commonSituations":"Trying to trace a whole tile or vector for debugging instead of a single element; forgetting to index the block first (arr[0, 0]).","solutions":["Trace a scalar element, e.g. trace_value(x[0, 0]) or trace_value(x.flatten()[0])","Trace a reduction of the tile, e.g. trace_value(x.sum()) if a representative value suffices (must be i32/f32)"],"exampleFix":"# before\ntrace_value(block, label='block')\n# after\ntrace_value(block[0, 0], label='block00')","handlingStrategy":"validation","validationCode":"assert not getattr(value, 'shape', None), 'trace_value needs a scalar'","typeGuard":"def is_scalar_aval(value) -> bool:\n    return not getattr(value, 'shape', None)","tryCatchPattern":null,"preventionTips":["Index a single element of the tile before tracing","Trace reductions (x.sum()) for tile-level insight"],"tags":["jax","pallas","tpu","debugging","scalar-required"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}