{"record":{"id":"6505cefb2b3c6790","repo":"jax-ml/jax","slug":"transformedref-must-have-been-indexed-before-passi","errorCode":null,"errorMessage":"TransformedRef must have been indexed before passing into jaxpr_call. Got {ref}.","messagePattern":"TransformedRef must have been indexed before passing into jaxpr_call\\. Got (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":1431,"sourceCode":"  the transformed references have been indexed.\n\n  Args:\n    jaxpr: The jaxpr to call.\n    *refs: The references to pass into the jaxpr.\n    program_ids: The loop-bound program IDs to pass into the jaxpr, or None if\n      the program ID corresponds to a parallel dimension.\n\n  Returns:\n    The outputs of the jaxpr.\n  \"\"\"\n  assert not jaxpr.outvars\n  flat_refs = []\n  ref_treedefs = []\n  ref: Any\n  for ref in refs:\n    if isinstance(ref, state_types.TransformedRef):\n      if not isinstance(ref.transforms[-1], indexing.NDIndexer):\n        raise ValueError(\n            \"TransformedRef must have been indexed before passing into\"\n            f\" jaxpr_call. Got {ref}.\"\n        )\n      ref = (ref.ref, ref.transforms)\n    flat_ref, treedef = tree_util.tree_flatten(ref)\n    flat_refs.extend(flat_ref)\n    ref_treedefs.append(treedef)\n  flat_program_ids, program_ids_treedef = tree_util.tree_flatten(program_ids)\n  return jaxpr_call_p.bind(\n      *flat_refs,\n      *flat_program_ids,\n      jaxpr=jaxpr,\n      ref_treedefs=tuple(ref_treedefs),\n      program_ids_treedef=program_ids_treedef,\n  )\n","sourceCodeStart":1413,"sourceCodeEnd":1447,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L1413-L1447","documentation":"When converting refs into jaxpr_call arguments, any TransformedRef must end with an NDIndexer transform (i.e., it must have been indexed). Passing an un-indexed TransformedRef (e.g., only transposed/sliced abstractly) is invalid.","triggerScenarios":"Building a jaxpr_call (used in ControlFlowInKernels / call-in-kernel) with a TransformedRef whose last transform is not an indexing.NDIndexer — e.g., a permutation or slice without a [...] index.","commonSituations":"Applying permute/swap_dims to a Ref and passing it to an in-kernel function call without indexing it first; internal Pallas control-flow plumbing errors.","solutions":["Index the ref before passing: ref[...] or a concrete index so the last transform is an NDIndexer","Pass the underlying ref plus explicit transforms/indexers instead of a bare transformed ref"],"exampleFix":"// before\ncall_in_kernel(f, ref_permuted)  # transformed but not indexed\n// after\ncall_in_kernel(f, ref_permuted[...])","handlingStrategy":"validation","validationCode":"from jax._src.pallas import indexing\nassert not isinstance(ref, state_types.TransformedRef) or isinstance(ref.transforms[-1], indexing.NDIndexer)","typeGuard":"def is_indexed_transformed_ref(ref) -> bool:\n    from jax._src import state_types\n    from jax._src.pallas import indexing\n    return (not isinstance(ref, state_types.TransformedRef)\n            or isinstance(ref.transforms[-1], indexing.NDIndexer))","tryCatchPattern":null,"preventionTips":["Index refs with [...] before passing to in-kernel calls","Avoid passing permuted/sliced refs without a final index"],"tags":["pallas","refs","jaxpr-call","validation","jax"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}