jax-ml/jax · error · ValueError

TransformedRef must have been indexed before passing into ja

Error message

TransformedRef must have been indexed before passing into jaxpr_call. Got {ref}.

What it means

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.

Source

Thrown at jax/_src/pallas/primitives.py:1431

  the transformed references have been indexed.

  Args:
    jaxpr: The jaxpr to call.
    *refs: The references to pass into the jaxpr.
    program_ids: The loop-bound program IDs to pass into the jaxpr, or None if
      the program ID corresponds to a parallel dimension.

  Returns:
    The outputs of the jaxpr.
  """
  assert not jaxpr.outvars
  flat_refs = []
  ref_treedefs = []
  ref: Any
  for ref in refs:
    if isinstance(ref, state_types.TransformedRef):
      if not isinstance(ref.transforms[-1], indexing.NDIndexer):
        raise ValueError(
            "TransformedRef must have been indexed before passing into"
            f" jaxpr_call. Got {ref}."
        )
      ref = (ref.ref, ref.transforms)
    flat_ref, treedef = tree_util.tree_flatten(ref)
    flat_refs.extend(flat_ref)
    ref_treedefs.append(treedef)
  flat_program_ids, program_ids_treedef = tree_util.tree_flatten(program_ids)
  return jaxpr_call_p.bind(
      *flat_refs,
      *flat_program_ids,
      jaxpr=jaxpr,
      ref_treedefs=tuple(ref_treedefs),
      program_ids_treedef=program_ids_treedef,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index the ref before passing: ref[...] or a concrete index so the last transform is an NDIndexer
2. Pass the underlying ref plus explicit transforms/indexers instead of a bare transformed ref

Example fix

// before
call_in_kernel(f, ref_permuted)  # transformed but not indexed
// after
call_in_kernel(f, ref_permuted[...])

from jax._src.pallas import indexing
assert not isinstance(ref, state_types.TransformedRef) or isinstance(ref.transforms[-1], indexing.NDIndexer)

def is_indexed_transformed_ref(ref) -> bool:
    from jax._src import state_types
    from jax._src.pallas import indexing
    return (not isinstance(ref, state_types.TransformedRef)
            or isinstance(ref.transforms[-1], indexing.NDIndexer))

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Sum of sizes {n} must be equal to dimension {axis} of the op
• Cannot pull iota along dimension {dimension} with None block
• Invalid memory space: {memory_space!r}
• dimension_semantics must be the same length as grid.
• Grid must consist of Python integers and JAX Arrays: {grid_t

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