{"record":{"id":"a934a18b3d18a4a5","repo":"jax-ml/jax","slug":"acc-ref-must-be-at-least-2d-got-shape-shape","errorCode":null,"errorMessage":"Acc ref must be at least 2D, got shape {shape}","messagePattern":"Acc ref must be at least 2D, got shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/core.py","lineNumber":215,"sourceCode":"        self, \"shape_invariant_numerics\", shape_invariant_numerics\n    )\n    object.__setattr__(self, \"use_tc_tiling_on_sc\", use_tc_tiling_on_sc)\n    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 \"","sourceCodeStart":197,"sourceCodeEnd":233,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/core.py#L197-L233","documentation":"Mosaic TPU accumulator refs (pltpu.ACC) map directly onto TPU MXU accumulator hardware, which is inherently 2D (sublanes x lanes). A ref whose shape has fewer than 2 dimensions cannot be backed by that hardware, so check_accumulator_ref raises ValueError immediately.","triggerScenarios":"Creating a Mosaic TPU kernel accumulator ref (e.g., via TensorCore compute tile plumbing or pltpu accumulator APIs) with a 1D or scalar shape, verified in __post_init__/__call__ of the wrapper class.","commonSituations":"Writing a TPU matmul kernel and allocating the accumulator with shape (N,) or () instead of (M, N); generalizing a GPU Pallas kernel whose accumulator was 1D.","solutions":["Give the accumulator at least 2 dimensions, e.g. reshape (K,) to (1, K) or use the intended (M, N) tile shape","Verify the minor dimension equals the MXU column size (see related checks) after fixing rank","If you only need a 1D result, compute into a 2D acc and reshape/squeeze after accumulator_load"],"exampleFix":"// before\nacc = pltpu.ACC  used with shape (256,)\n// after\nacc_shape = (1, 256)  # 2D accumulator; reshape after load\nacc = ... allocate with acc_shape","handlingStrategy":"validation","validationCode":"def valid_acc_shape(shape):\n    return len(shape) >= 2","typeGuard":"def is_2d_plus(shape: tuple[int, ...]) -> TypeGuard[tuple[int, int, ...]]: return len(shape) >= 2","tryCatchPattern":null,"preventionTips":["Always allocate ACC refs as (M, N) tiles","Centralize accumulator allocation in one helper that asserts rank >= 2"],"tags":["jax","pallas","tpu","mosaic","accumulator","shape"],"backgroundTag":"invalid-shape-for-hardware","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}