jax-ml/jax · error · ValueError

Array shape mismatch: expected {ty.shape} got {v.shape}.

Error message

Array shape mismatch: expected {ty.shape} got {v.shape}.

What it means

Same per-argument type check as the dtype case, but for shapes: the FragmentedArray's runtime shape must equal the shape declared in the ShapeDtypeStruct passed to inline_mgpu.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3533

@discharge.register_discharge_rule(inline_mgpu_p)
def _inline_mgpu_discharge(*args, **kwargs):
  del args, kwargs
  raise NotImplementedError("inline_mgpu_p does not support discharge.")


def _type_check_mgpu_lane_semantics(v, ty):
  match (ty, v):
    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):
      pass
    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):
      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)
      if v.mlir_dtype != mlir_dtype:
        raise ValueError(
            f"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}."
        )
      if ty.shape != v.shape:
        raise ValueError(
            f"Array shape mismatch: expected {ty.shape} got {v.shape}."
        )
      if v.layout != ty.layout.to_mgpu():
        raise ValueError(
            f"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}."
        )
    case (SomeLayout(), mgpu.FragmentedArray()):
      if ty.to_mgpu() != v.layout:
        raise ValueError(f"Unexpected layout for {v} (expected: {ty})")
    case _:
      raise ValueError(f"Unexpected type {ty} for value {v}")


def _inline_mgpu_flat_transformed_args(
    ctx: lowering.LoweringRuleContext,
    flat_args_and_transforms,
    flat_arg_types,
    pytree_args,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Update the declared shape to the actual produced shape
2. Reshape/splat the output inside the wrapped function to the declared shape

Example fix

# before
return_type=ShapeDtypeStruct((16,16), dt)
# after
return_type=ShapeDtypeStruct((16,32), dt)  # matches actual output

assert v.shape == declared.shape

When it happens

Trigger: Declaring ShapeDtypeStruct((16,16), dtype) while the wrapped function emits a FragmentedArray of shape (16,32), or returning a scalar where an array was declared.

Common situations: Hard-coding block shapes that later change when block specs are updated.

Related errors

• Block shape for {origin} (= {block_shape}) must have the sam
• Stacking only supported when the block size along the stack
• Sum of sizes {n} must be equal to dimension {axis} of the op
• Block size must be a multiple of the input size. Got block {
• Every block dimension must be either a multiple or factor of

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