jax-ml/jax · error · ValueError

Size mismatch for group {group}: expected {shape_val}, got {

Error message

Size mismatch for group {group}: expected {shape_val}, got {known_product}

What it means

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.

Source

Thrown at jax/_src/pallas/einshape.py:198

            f"Inconsistent size for {name}: {dim_sizes[name]} vs {shape_val}"
        )
      dim_sizes[name] = shape_val
    else:
      # We have a merged dimension on LHS, need to split
      known_product = 1
      unknown_dims = []
      for name in group:
        if name in sizes:
          dim_sizes[name] = sizes[name]
          known_product *= sizes[name]
        elif name in dim_sizes:
          known_product *= dim_sizes[name]
        else:
          unknown_dims.append(name)

      if not unknown_dims:
        if known_product != shape_val:
          raise ValueError(
              f"Size mismatch for group {group}: expected {shape_val}, got"
              f" {known_product}"
          )
      elif len(unknown_dims) == 1:
        if shape_val % known_product != 0:
          raise ValueError(
              f"Cannot split size {shape_val} with known sizes {known_product}"
          )
        inferred_size = shape_val // known_product
        dim_sizes[unknown_dims[0]] = inferred_size
      else:
        raise ValueError(
            f"Ambiguous split for {group} with size {shape_val}. Unknowns:"
            f" {unknown_dims}. Provide sizes via kwargs."
        )
  return dim_sizes

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Adjust the group members so their sizes multiply exactly to the axis extent (e.g. 224 -> (16 14))
2. If exactly one member size is unknown, einshape can infer it — leave one letter's size unspecified rather than guessing wrong values
3. Add an assert prod(group_dims) == axis_size in your code before calling einshape so failures point at your constants

Example fix

# before
t = get_einshape_transforms('(h w) c -> h w c', (224, 3))  # h*w unknown/wrong -> mismatch

# after
t = get_einshape_transforms('(h w) c -> h w c', (224, 3))
# ensure prior dims: dim_sizes['h']=14, dim_sizes['w']=16 defined elsewhere, 14*16 == 224

import math
def group_products_ok(parsed_side, dim_sizes, shape):
    for group, s in zip(parsed_side, shape):
        if len(group) > 1:
            known = [dim_sizes[n] for n in group if n in dim_sizes]
            if len(known) == len(group) and math.prod(known) != s:
                return False
    return True

null

try:
    t = get_einshape_transforms(eq, shape)
except ValueError as e:
    if 'Size mismatch for group' in str(e):
        raise ValueError(f'grouped dims in {eq!r} do not multiply to axis sizes {shape}') from None
    raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Unmatched parenthesis in {s!r}
• Equation must contain exactly one '->'
• Inconsistent size for {name}: {dim_sizes[name]} vs {shape_va
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3f652ed8d6dc6de8. Report an issue: GitHub.