jax-ml/jax · error · ValueError

Array shape mismatch: expected

Error message

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

What it means

Runtime consistency check when lowering run_scoped bodies with Lane semantics: a mgpu.FragmentedArray produced for a value declared as a ShapeDtypeStruct must match the declared shape. A mismatch means the value computed inside the scoped block doesn't conform to the output type the enclosing primitive declared.

Solutions

Update the declared shape to the actual produced shape
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

Defensive patterns

Strategy: validation

Validate before calling

assert v.shape == declared.shape

Prevention

Derive declared shapes from block specs rather than hard-coding
Assert output shapes in kernel unit tests

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

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

Appendix: 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)