jax-ml/jax · error · ValueError

Mismatched type: {a, t}

Error message

Mismatched type: {a, t}

What it means

Each runtime argument passed to an inline_mgpu-wrapped function is matched against its declared type: refs with transforms, plain state.AbstractRef for RefType, and FragmentedArray for layout types. If the runtime value's aval does not match the declared type entry, this error is raised with the (value, type) pair.

Source

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

      if treedef != treedef_ty:
        raise ValueError(f"Mismatched type shape: {treedef} != {treedef_ty}")

      # Strip the transforms from the refs since they will be recorded in
      # the types.
      ref_transforms: list[Any] = []
      raw_flat_args = []
      for a, t in zip(flat_args, flat_arg_types):
        if isinstance(a, state_types.TransformedRef) and isinstance(t, RefType):
          raw_flat_args.append(a.ref)
          ref_transforms.append(a.transforms)
        elif isinstance(aval := jax_core.typeof(a), jax_core.ShapedArray) and isinstance(t, SomeLayout):
          raw_flat_args.append(a)
          ref_transforms.append(None)
        elif isinstance(aval, state.AbstractRef) and isinstance(t, RefType):
          raw_flat_args.append(a)
          ref_transforms.append(())
        else:
          raise ValueError(f"Mismatched type: {a, t}")

      flat_ref_transforms, pytree_ref_transforms = jax.tree.flatten(ref_transforms)
      flat_ret = inline_mgpu_p.bind(
          *raw_flat_args,
          *flat_ref_transforms,
          flat_arg_types=tuple(flat_arg_types),
          flat_ret_ty=tuple(flat_ret_ty),
          pytree_ret_ty=pytree_ret_ty,
          pytree_args=treedef,
          pytree_ref_transforms=pytree_ref_transforms,
          mgpu_fn=f,
      )
      return jax.tree.unflatten(pytree_ret_ty, flat_ret)
    return wrapper

  return inner

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder/convert arguments so each matches the corresponding entry in arg_types
2. Wrap scalars as FragmentedArray (e.g. via mgpu.splat / _array_splat) and buffers as refs

Example fix

# before
f = inline_mgpu(impl, arg_types=[RefType(...)], ...)
f(3.0)
# after
f = inline_mgpu(impl, arg_types=[Layout(...)], ...)
f(splatted_scalar)

for a, t in zip(flat_args, flat_arg_types):
    assert matches(a, t), (a, t)

def arg_matches(a, t):
    from jax._src import state
    return (isinstance(a, state.AbstractRef) and isinstance(t, RefType)) or (hasattr(a, 'layout') and isinstance(t, SomeLayout))

When it happens

Trigger: Passing a python scalar or a plain array where a RefType was declared, or passing a FragmentedArray where a different layout type was declared.

Common situations: Mixed up ordering of arguments relative to arg_types; passing transformed refs where plain refs expected.

Related errors

• with_memory_space_constraint only supports arrays.
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
• inline_mgpu_p only supports only SomeLayout and RefType arg
• Mismatched type shape: {treedef} != {treedef_ty}
• inline_mgpu_p does not support discharge.

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