{"record":{"id":"68fbab4bdb4ddffc","repo":"jax-ml/jax","slug":"mismatched-type-a-t","errorCode":null,"errorMessage":"Mismatched type: {a, t}","messagePattern":"Mismatched type: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3470,"sourceCode":"      if treedef != treedef_ty:\n        raise ValueError(f\"Mismatched type shape: {treedef} != {treedef_ty}\")\n\n      # Strip the transforms from the refs since they will be recorded in\n      # the types.\n      ref_transforms: list[Any] = []\n      raw_flat_args = []\n      for a, t in zip(flat_args, flat_arg_types):\n        if isinstance(a, state_types.TransformedRef) and isinstance(t, RefType):\n          raw_flat_args.append(a.ref)\n          ref_transforms.append(a.transforms)\n        elif isinstance(aval := jax_core.typeof(a), jax_core.ShapedArray) and isinstance(t, SomeLayout):\n          raw_flat_args.append(a)\n          ref_transforms.append(None)\n        elif isinstance(aval, state.AbstractRef) and isinstance(t, RefType):\n          raw_flat_args.append(a)\n          ref_transforms.append(())\n        else:\n          raise ValueError(f\"Mismatched type: {a, t}\")\n\n      flat_ref_transforms, pytree_ref_transforms = jax.tree.flatten(ref_transforms)\n      flat_ret = inline_mgpu_p.bind(\n          *raw_flat_args,\n          *flat_ref_transforms,\n          flat_arg_types=tuple(flat_arg_types),\n          flat_ret_ty=tuple(flat_ret_ty),\n          pytree_ret_ty=pytree_ret_ty,\n          pytree_args=treedef,\n          pytree_ref_transforms=pytree_ref_transforms,\n          mgpu_fn=f,\n      )\n      return jax.tree.unflatten(pytree_ret_ty, flat_ret)\n    return wrapper\n\n  return inner\n\n","sourceCodeStart":3452,"sourceCodeEnd":3488,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3452-L3488","documentation":"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.","triggerScenarios":"Passing a python scalar or a plain array where a RefType was declared, or passing a FragmentedArray where a different layout type was declared.","commonSituations":"Mixed up ordering of arguments relative to arg_types; passing transformed refs where plain refs expected.","solutions":["Reorder/convert arguments so each matches the corresponding entry in arg_types","Wrap scalars as FragmentedArray (e.g. via mgpu.splat / _array_splat) and buffers as refs"],"exampleFix":"# before\nf = inline_mgpu(impl, arg_types=[RefType(...)], ...)\nf(3.0)\n# after\nf = inline_mgpu(impl, arg_types=[Layout(...)], ...)\nf(splatted_scalar)","handlingStrategy":"type-guard","validationCode":"for a, t in zip(flat_args, flat_arg_types):\n    assert matches(a, t), (a, t)","typeGuard":"def arg_matches(a, t):\n    from jax._src import state\n    return (isinstance(a, state.AbstractRef) and isinstance(t, RefType)) or (hasattr(a, 'layout') and isinstance(t, SomeLayout))","tryCatchPattern":null,"preventionTips":["Keep argument order in lockstep with arg_types","Splat scalars before passing"],"tags":["jax","pallas","inline-mgpu","type-mismatch"],"backgroundTag":"argument-type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}