{"record":{"id":"4ad161ad9e5ab8b7","repo":"jax-ml/jax","slug":"mismatched-type-shape-treedef-treedef-ty","errorCode":null,"errorMessage":"Mismatched type shape: {treedef} != {treedef_ty}","messagePattern":"Mismatched type shape: (.+?) != (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3453,"sourceCode":"      :class:`~jax.experimental.pallas.mosaic_gpu.ShapeDtypeStruct`\\s\n      representing the arrays returned by the decorated function.\n  \"\"\"\n  flat_arg_types, treedef_ty = jax.tree.flatten(tuple(arg_types))\n  flat_ret_ty, pytree_ret_ty = jax.tree.flatten(return_type)\n  if return_type and not all(isinstance(r, ShapeDtypeStruct) for r in flat_ret_ty):\n    raise ValueError(\n        \"inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types.\"\n    )\n  if not all(isinstance(r, (SomeLayout, RefType)) for r in flat_arg_types):\n    raise ValueError(\n        \"inline_mgpu_p only supports only SomeLayout and RefType arg types.\"\n    )\n\n  def inner(f):\n    def wrapper(*args):\n      flat_args, treedef = jax.tree.flatten(tuple(args))\n      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","sourceCodeStart":3435,"sourceCodeEnd":3471,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3435-L3471","documentation":"The function wrapped by inline_mgpu must be called with arguments whose pytree structure exactly matches the arg_types structure passed at decoration time. The wrapper flattens both and compares treedefs; any difference (extra/missing args, different nesting) raises this error.","triggerScenarios":"Decorating with arg_types=[RefType, RefType] but calling f with one argument, or passing args nested differently ((a,b),c vs a,b,c).","commonSituations":"Refactoring the wrapped function signature after decoration; passing optional args conditionally.","solutions":["Call the wrapped function with exactly the same arg structure as arg_types","Update arg_types to match the new call signature"],"exampleFix":"# before\nf = inline_mgpu(impl, arg_types=[ty_a, ty_b], return_type=...)\nf(a)\n# after\nf(a, b)","handlingStrategy":"validation","validationCode":"assert jax.tree.structure(args) == jax.tree.structure(arg_types_tuple)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Call the wrapped function with a fixed signature","Regenerate arg_types whenever the signature changes"],"tags":["jax","pallas","inline-mgpu","pytree","signature-mismatch"],"backgroundTag":"argument-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}