jax-ml/jax · error · ValueError

Expected {result_shape} but got {fa.shape}

Error message

Expected {result_shape} but got {fa.shape}

What it means

The shape of the FragmentedArray returned from an inline_mgpu callback must equal the shape of the declared vector result type. A mismatch raises ValueError with both shapes.

Source

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

        fa = mgpu.FragmentedArray.splat(arg, (), is_signed=is_signed)
        fn_inputs.append(fa)

    args = jax.tree.unflatten(pytree_args, fn_inputs)
    inner_ret = mgpu_fn(ctx.launch_ctx, *args)
    if inner_ret is None:
      inner_ret = []
    elif not isinstance(inner_ret, tuple) and not isinstance(inner_ret, list):
      inner_ret = [inner_ret]
    ir_ret = []
    for fa, result_ty, out_layout in zip(
        inner_ret, results_ty, out_layouts, strict=True
    ):
      if not isinstance(fa, mgpu.FragmentedArray):
        raise ValueError(f"Expected a FragmentedArray, but got: {fa}")
      if isinstance(result_ty, ir.VectorType):
        result_shape = ir.VectorType(result_ty).shape
        if fa.shape != tuple(result_shape):
          raise ValueError(f"Expected {result_shape} but got {fa.shape}")
        if out_layout != mgpu.layouts.to_layout_attr(fa.layout):
          raise ValueError(
              f"Output layout {out_layout} does not match the layout of the"
              f" returned fragmented array {fa.layout}."
          )
        ir_ret.append(
            mgpu.dialect_lowering.fragmented_array_to_ir(fa, result_ty)
        )
      else:  # scalar case.
        assert out_layout is None
        if fa.shape:
          raise ValueError(f"Expected 0D shape, but got {fa.shape}")
        if not isinstance(fa.layout, mgpu.WGSplatFragLayout):
          raise ValueError(f"Expected WGSplatFragLayout, but got {fa.layout}")
        value = fa.registers.item()
        ir_ret.append(value)

    mgpu.dialect.return_(ir_ret)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the declared return aval shape match exactly what the callback returns
2. Remove implicit reshapes or broadcasts in the callback, or declare the broadcast shape explicitly

Example fix

// before
ret_ty = jax.ShapeDtypeStruct((16, 16), jnp.float32)  # callback returns (8,8)
// after
ret_ty = jax.ShapeDtypeStruct((8, 8), jnp.float32)

assert tuple(fa.shape) == tuple(ret_ty.shape), (fa.shape, ret_ty.shape)

When it happens

Trigger: Callback returns a FragmentedArray whose .shape differs from ir.VectorType(result_ty).shape for that output.

Common situations: Declaring output avals with a different shape than what the callback computes; broadcasting applied inside the callback changing shapes.

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/d00c3ead6f0912c8. Report an issue: GitHub.