jax-ml/jax · error · ValueError

Unexpected type for value

Error message

Unexpected type {ty} for value {v}

What it means

The inline_mgpu lane-semantics type checker only understands (RefType, MemRef value), (ShapeDtypeStruct, FragmentedArray) and (SomeLayout, FragmentedArray) pairs. Any other (type, value) combination falls through to this generic 'Unexpected type' error.

Solutions

1. Convert the value to a supported runtime representation (splat scalars into FragmentedArray via mgpu.splat)
2. Use only RefType / ShapeDtypeStruct / layout types in arg_types and return_type

Example fix

# before
f(True)
# after
f(mgpu.splat(True, layout))  # or use mgpu.c for MLIR constants inside the impl

def supported_pair(ty, v):
    return (isinstance(ty, RefType) and isinstance(v.type, ir.MemRefType)) or (isinstance(ty, (ShapeDtypeStruct, SomeLayout)) and isinstance(v, mgpu.FragmentedArray))
assert supported_pair(ty, v)

def supported_pair(ty, v):
    return (isinstance(ty, RefType) and isinstance(v.type, ir.MemRefType)) or (isinstance(ty, (ShapeDtypeStruct, SomeLayout)) and isinstance(v, mgpu.FragmentedArray))

When it happens

Trigger: Declaring a type entry that is neither RefType nor ShapeDtypeStruct nor SomeLayout-derived, or passing a runtime value that is not an ir.Value memref or FragmentedArray (e.g. a python scalar, bool).

Common situations: Passing predicates as plain python bools instead of FragmentedArray splats; custom type annotations in arg_types.

Related errors

• inline_mgpu_p only supports only SomeLayout and RefType arg…
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return…
• Array dtype mismatch: expected
• Array layout mismatch: expected
• Array shape mismatch: expected

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

Appendix: source

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

    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,
    pytree_ref_transforms,
  ) -> Sequence[ir.Value | mgpu.FragmentedArray]:
  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]
  flat_arg_avals = ctx.avals_in[:pytree_args.num_leaves]
  ref_transforms = pytree_ref_transforms.unflatten(flat_args_and_transforms[pytree_args.num_leaves:])
  ref_transform_avals = pytree_ref_transforms.unflatten(ctx.avals_in[pytree_args.num_leaves:])
  is_wg_semantics = (
      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
  )
  is_warp_semantics = (
      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp

View on GitHub (pinned to 1e1c6a8fc0)