{"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":"The accumulator's last (minor) dimension must exactly equal the MXU column size reported by TPU info (typically the lane width). Any other minor dimension cannot be laid out in accumulator hardware.","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","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}