jax-ml/jax · error · ValueError

Unexpected type {ty} for value {v}

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.

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)

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 plgpu.ShapeDtypeStruct return ty
• inline_mgpu_p only supports only SomeLayout and RefType arg
• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block dimension type: {type(bd)}

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