{"record":{"id":"a5e0e449d75a106e","repo":"jax-ml/jax","slug":"the-minor-dimension-size-of-an-accumulator-ref-mus","errorCode":null,"errorMessage":"The minor dimension size of an accumulator ref must be {info.mxu_column_size} but got {n}","messagePattern":"The minor dimension size of an accumulator ref must be (.+?) but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/core.py","lineNumber":232,"sourceCode":"  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:\n    raise ValueError(\n        f\"The product of the major dimensions must be a multiple of \"\n        f\"{info.num_sublanes}, but got {m}\"\n    )\n\n\nclass MemoryRef(pallas_core.MemoryRef):\n\n  def __matmul__(self, other, /):\n    if not isinstance(other, pallas_core.Mesh):\n      return NotImplemented\n    return dataclasses.replace(self, memory_space=self.memory_space @ other)\n\n","sourceCodeStart":214,"sourceCodeEnd":250,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/core.py#L214-L250","documentation":"Validation in check_accumulator_ref: an accumulator ref's minor (last) dimension must equal the TPU MXU column size (info.mxu_column_size), since accumulators live in MXU columns; a ref whose last dimension n differs cannot be mapped onto the hardware accumulator and raises ValueError.","triggerScenarios":"Allocating an ACC ref whose shape[-1] differs from info.mxu_column_size, e.g. shape (128, 100) on hardware with 128-wide MXU columns requiring exactly 128.","commonSituations":"Using non-power-of-two or arbitrary N dimensions in matmul tiles; porting kernels between TPU generations with different MXU widths.","solutions":["Set shape[-1] to exactly info.mxu_column_size (query it at runtime)","Pad the N dimension of your GEMM tile up to the MXU column size and slice after loading the accumulator","Regenerate tile shapes per target TPU generation instead of hardcoding"],"exampleFix":"// before\nacc = alloc_acc((M, 100))\n// after\ninfo = tpu_info.get_tpu_info()\nacc = alloc_acc((M, info.mxu_column_size))\nresult = acc.load()[:, :100]","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic import tpu_info\ninfo = tpu_info.get_tpu_info()\nassert shape[-1] == info.mxu_column_size, (shape, info.mxu_column_size)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive tile N from info.mxu_column_size at runtime","Pad N and slice after load when the math needs a smaller N"],"tags":["jax","pallas","tpu","accumulator","shape"],"backgroundTag":"invalid-shape-for-hardware","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}