{"record":{"id":"77c28fea7782f366","repo":"jax-ml/jax","slug":"acc-ref-dtype-must-be-float32-or-int32-got-dtype","errorCode":null,"errorMessage":"Acc ref dtype must be float32 or int32, got {dtype}","messagePattern":"Acc ref dtype must be float32 or int32, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/core.py","lineNumber":218,"sourceCode":"    object.__setattr__(self, \"needs_layout_passes\", needs_layout_passes)\n    object.__setattr__(\n        self,\n        \"fuse_transposed_lhs_in_matmul\",\n        fuse_transposed_lhs_in_matmul,\n    )\n    object.__setattr__(self, \"opt_level\", opt_level)\n\n  # Replace is a method, not a field.\n  replace = dataclasses.replace\n\n\ndef check_accumulator_ref(shape: tuple[int, ...], dtype: jnp.dtype, mxu_id: int):\n  from jax._src.pallas.mosaic import tpu_info  # pyrefly: ignore[missing-module-attribute]\n  if len(shape) < 2:\n    raise ValueError(f\"Acc ref must be at least 2D, got shape {shape}\")\n\n  if dtype not in (jnp.float32, jnp.int32):\n    raise ValueError(\n        f\"Acc ref dtype must be float32 or int32, got {dtype}\")\n\n  info = tpu_info.get_tpu_info()\n  if not info.num_accumulators:\n    raise ValueError(\n        f\"Accumulators are not available on TPU {info.chip_version}\"\n    )\n\n  if mxu_id < 0 or mxu_id >= info.num_mxus:\n    raise ValueError(f\"mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}\")\n\n  m, n = math.prod(shape[:-1]), shape[-1]\n  if n != info.mxu_column_size:\n    raise ValueError(\n        f\"The minor dimension size of an accumulator ref must be \"\n        f\"{info.mxu_column_size} but got {n}\"\n    )\n  if m <= 0 or m % info.num_sublanes != 0:","sourceCodeStart":200,"sourceCodeEnd":236,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/core.py#L200-L236","documentation":"TPU MXU accumulators only support float32 and int32 element types. Passing an accumulator ref with any other dtype (bfloat16, float16, float64, int8, ...) fails this validation before any compilation happens.","triggerScenarios":"Allocating a Mosaic TPU accumulator with jnp.bfloat16 or another dtype instead of jnp.float32/jnp.int32 in check_accumulator_ref (called during wrapper __call__/__post_init__).","commonSituations":"Writing bf16 matmul kernels (common on TPU) and assuming the accumulator can also be bf16; converting GPU Pallas kernels that accumulate in fp16.","solutions":["Declare the accumulator as jnp.float32 (standard for matmul accumulation) or jnp.int32","Cast to bfloat16 only after loading the accumulator out to VMEM","Double-check any dtype inference (e.g., from jnp.zeros default) that might produce non-float32"],"exampleFix":"// before\nacc = kernel_init(dtype=jnp.bfloat16)\n// after\nacc = kernel_init(dtype=jnp.float32)\nout = acc.load().astype(jnp.bfloat16)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert acc_dtype in (jnp.float32, jnp.int32), acc_dtype","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default accumulators to float32","Cast to bf16 only after acc.load()"],"tags":["jax","pallas","tpu","accumulator","dtype"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}