{"record":{"id":"3f652ed8d6dc6de8","repo":"jax-ml/jax","slug":"size-mismatch-for-group-group-expected-shape-v","errorCode":null,"errorMessage":"Size mismatch for group {group}: expected {shape_val}, got {known_product}","messagePattern":"Size mismatch for group (.+?): expected (.+?), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/einshape.py","lineNumber":198,"sourceCode":"            f\"Inconsistent size for {name}: {dim_sizes[name]} vs {shape_val}\"\n        )\n      dim_sizes[name] = shape_val\n    else:\n      # We have a merged dimension on LHS, need to split\n      known_product = 1\n      unknown_dims = []\n      for name in group:\n        if name in sizes:\n          dim_sizes[name] = sizes[name]\n          known_product *= sizes[name]\n        elif name in dim_sizes:\n          known_product *= dim_sizes[name]\n        else:\n          unknown_dims.append(name)\n\n      if not unknown_dims:\n        if known_product != shape_val:\n          raise ValueError(\n              f\"Size mismatch for group {group}: expected {shape_val}, got\"\n              f\" {known_product}\"\n          )\n      elif len(unknown_dims) == 1:\n        if shape_val % known_product != 0:\n          raise ValueError(\n              f\"Cannot split size {shape_val} with known sizes {known_product}\"\n          )\n        inferred_size = shape_val // known_product\n        dim_sizes[unknown_dims[0]] = inferred_size\n      else:\n        raise ValueError(\n            f\"Ambiguous split for {group} with size {shape_val}. Unknowns:\"\n            f\" {unknown_dims}. Provide sizes via kwargs.\"\n        )\n  return dim_sizes\n\n","sourceCodeStart":180,"sourceCodeEnd":216,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/einshape.py#L180-L216","documentation":"For a parenthesized dimension group like '(abc)', einshape multiplies the known member sizes and requires the product to equal the corresponding input axis extent. When all members' sizes are known and the product differs from the axis size, _get_einshape_dims raises ValueError 'Size mismatch for group'. Essentially the grouped-reshape factorization is arithmetically wrong.","triggerScenarios":"Equation like 'a(bc) -> abc' where a*? ... specifically a group whose known sizes multiply to something other than the axis extent, e.g. get_einshape_transforms('(ab) -> ab', shape=(2, 6)) — but here with all dims known, e.g. '(ab)(cd) -> abcd' where a=2,b=3 but axis0=7. Typical direct case: get_einshape_transforms('(h w) c -> hwc', shape=(224, 3)) when h*w != 224.","commonSituations":"Hardcoding spatial factors (h, w) that no longer multiply to the flattened axis after resizing inputs; off-by-one in grid dims; porting reshape(224*224) style code to einshape with wrong factorization.","solutions":["Adjust the group members so their sizes multiply exactly to the axis extent (e.g. 224 -> (16 14))","If exactly one member size is unknown, einshape can infer it — leave one letter's size unspecified rather than guessing wrong values","Add an assert prod(group_dims) == axis_size in your code before calling einshape so failures point at your constants"],"exampleFix":"# before\nt = get_einshape_transforms('(h w) c -> h w c', (224, 3))  # h*w unknown/wrong -> mismatch\n\n# after\nt = get_einshape_transforms('(h w) c -> h w c', (224, 3))\n# ensure prior dims: dim_sizes['h']=14, dim_sizes['w']=16 defined elsewhere, 14*16 == 224","handlingStrategy":"validation","validationCode":"import math\ndef group_products_ok(parsed_side, dim_sizes, shape):\n    for group, s in zip(parsed_side, shape):\n        if len(group) > 1:\n            known = [dim_sizes[n] for n in group if n in dim_sizes]\n            if len(known) == len(group) and math.prod(known) != s:\n                return False\n    return True","typeGuard":"null","tryCatchPattern":"try:\n    t = get_einshape_transforms(eq, shape)\nexcept ValueError as e:\n    if 'Size mismatch for group' in str(e):\n        raise ValueError(f'grouped dims in {eq!r} do not multiply to axis sizes {shape}') from None\n    raise","preventionTips":["Verify prod(group member sizes) equals each axis extent before calling einshape","Derive factor pairs from the actual tensor shape (e.g. factor 224 into 16x14) rather than hardcoding","Leave one member of a group size-unknown when possible so einshape infers it correctly"],"tags":["jax","einshape","einsum-notation","shape-mismatch","grouped-reshape"],"backgroundTag":"einsum-dimension-size-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}