{"record":{"id":"09172f41b94f161c","repo":"jax-ml/jax","slug":"the-product-of-the-major-dimensions-must-be-a-mult","errorCode":null,"errorMessage":"The product of the major dimensions must be a multiple of {info.num_sublanes}, but got {m}","messagePattern":"The product of the major dimensions must be a multiple of (.+?), but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/core.py","lineNumber":237,"sourceCode":"        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\nclass MemorySpace(enum.Enum):\n  VMEM = \"vmem\"\n  VMEM_SHARED = \"vmem_shared\"\n  SMEM = \"smem\"\n  CMEM = \"cmem\"","sourceCodeStart":219,"sourceCodeEnd":255,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/core.py#L219-L255","documentation":"The product of the accumulator's major dimensions (all dims except the last) must be a positive multiple of num_sublanes, because accumulator rows are distributed across TPU sublanes. Zero or non-multiple totals are rejected.","triggerScenarios":"Allocating an ACC ref where prod(shape[:-1]) is 0 or not divisible by info.num_sublanes, e.g., major dim 100 when num_sublanes is 128.","commonSituations":"Choosing M tile sizes that don't match sublane counts; ragged batch dims flattened into the accumulator.","solutions":["Round the product of major dims up to the next multiple of num_sublanes (pad M)","Ensure M > 0 (e.g., a leading dimension of 0 from an empty batch)","Query info.num_sublanes and derive tile sizes from it per target chip"],"exampleFix":"// before\nacc = alloc_acc((100, N))\n// after\ninfo = tpu_info.get_tpu_info()\nM = ceil(100 / info.num_sublanes) * info.num_sublanes\nacc = alloc_acc((M, N))","handlingStrategy":"validation","validationCode":"import math\nfrom jax._src.pallas.mosaic import tpu_info\ninfo = tpu_info.get_tpu_info()\nm = math.prod(shape[:-1])\nassert m > 0 and m % info.num_sublanes == 0, (m, info.num_sublanes)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Round major-dim products up to num_sublanes multiples","Guard against zero-sized leading dims from empty batches"],"tags":["jax","pallas","tpu","accumulator","shape","alignment"],"backgroundTag":"invalid-shape-for-hardware","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}