jax-ml/jax · error · ValueError

Expected a FragmentedArray, but got: {fa}

Error message

Expected a FragmentedArray, but got: {fa}

What it means

Inside the inline_mgpu lowering, each returned value must be an mgpu.FragmentedArray when the corresponding result type is a vector type. Getting a plain array or other object raises this ValueError.

Source

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

        )
        fn_inputs.append(fa)
      else:  # scalar case.
        is_signed = mgpu_utils.is_signed(aval.dtype)
        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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the callback produces FragmentedArray outputs (e.g. via mgpu ops that yield fragments)
2. If returning a scalar, declare the return type as scalar so the scalar path is taken

from jax._src.pallas.mosaic_gpu import mgpu
assert all(isinstance(v, mgpu.FragmentedArray) for v in rets)

def is_fragmented(v): return isinstance(v, mgpu.FragmentedArray)

When it happens

Trigger: An inline_mgpu callback under warp-group semantics returns a raw value (numpy array, scalar, etc.) where a FragmentedArray is expected by the lowering.

Common situations: Returning constants, host-side arrays, or scalars in vector-typed output slots; mismatch between declared vector return type and scalar-producing callback body.

Related errors

• Unsupported core type: {core_type}
• Too many dynamic shapes in the input. Mosaic currently only
• program id was requested but no grid was provided.
• Invalid axis {axis} for num_programs
• out_dtype argument in binary_op_lowering_rule_wg

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