jax-ml/jax · error · ValueError

Multi-ref TransformedRef requires a single transform: {self}

Error message

Multi-ref TransformedRef requires a single transform: {self}

What it means

TransformedRef.__post_init__ validates that a multiref (a TransformedRef wrapping a tuple of refs) carries exactly one transform. Composing multiple transforms on a multiref is not representable, so a ValueError is raised.

Source

Thrown at jax/_src/state/types.py:296

    from jax._src.state import indexing
    indexer = indexing.NDIndexer.from_indices_shape(slc, self.ref_or_view.shape)
    if (
        isinstance(self.ref_or_view, TransformedRef)
        and not self.ref_or_view.multiref
    ):
      view = self.ref_or_view
      return TransformedRef(view.ref, (*view.transforms, indexer))
    return TransformedRef(self.ref_or_view, (indexer,))


@dataclasses.dataclass(frozen=True)
class TransformedRef:
  ref: Any
  transforms: tuple[Transform, ...]

  def __post_init__(self):
    if self.multiref and len(self.transforms) != 1:
      raise ValueError(
          f"Multi-ref TransformedRef requires a single transform: {self}"
      )
    if any(isinstance(t, MultiRefTransform) for t in self.transforms):
      assert self.multiref and len(self.transforms) == 1

  @property
  def multiref(self) -> bool:
    if isinstance(self.ref, Sequence):
      if all(isinstance(x, int) for x in self.ref):
        return False  # self.ref is an array's shape. This happens in lowering.
      return True
    return False

  @property
  def is_dynamic_size(self):
    return any(not isinstance(i, int) for i in self.shape)

  @functools.cached_property

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply the transform to each individual ref instead of the group
2. Flatten to a single ref before chaining transforms
3. Restructure so the multiref only ever receives one transform

Example fix

// before
group.transpose(perm).reshape(new_shape)
// after
for r in group_ref.ref:
    r.transpose(perm).reshape(new_shape)

if is_multiref(ref) and ref.transforms:
    raise ValueError("apply to individual refs instead")

def is_multiref(ref):
    return isinstance(getattr(ref, "ref", None), tuple)

When it happens

Trigger: Calling chained methods like ref_group.transpose(...).reshape(...) or .bitcast(...) then another transform on a multiref.

Common situations: Trying to apply .transpose to a multiref, or bitcast+reshape chains on grouped refs.

Related errors

• Cannot select from Refs of different types: {types}
• Cannot resolve attribute {name} from: {attrs}
• Transpose with multiref is not supported.
• Found inconsistent memory spaces in multiref: {self.ref}
• numpy masked arrays are not supported as direct inputs to JA

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