{"record":{"id":"ebeed3118a8b585b","repo":"jax-ml/jax","slug":"multi-ref-transformedref-requires-a-single-transfo","errorCode":null,"errorMessage":"Multi-ref TransformedRef requires a single transform: {self}","messagePattern":"Multi-ref TransformedRef requires a single transform: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/types.py","lineNumber":296,"sourceCode":"    from jax._src.state import indexing\n    indexer = indexing.NDIndexer.from_indices_shape(slc, self.ref_or_view.shape)\n    if (\n        isinstance(self.ref_or_view, TransformedRef)\n        and not self.ref_or_view.multiref\n    ):\n      view = self.ref_or_view\n      return TransformedRef(view.ref, (*view.transforms, indexer))\n    return TransformedRef(self.ref_or_view, (indexer,))\n\n\n@dataclasses.dataclass(frozen=True)\nclass TransformedRef:\n  ref: Any\n  transforms: tuple[Transform, ...]\n\n  def __post_init__(self):\n    if self.multiref and len(self.transforms) != 1:\n      raise ValueError(\n          f\"Multi-ref TransformedRef requires a single transform: {self}\"\n      )\n    if any(isinstance(t, MultiRefTransform) for t in self.transforms):\n      assert self.multiref and len(self.transforms) == 1\n\n  @property\n  def multiref(self) -> bool:\n    if isinstance(self.ref, Sequence):\n      if all(isinstance(x, int) for x in self.ref):\n        return False  # self.ref is an array's shape. This happens in lowering.\n      return True\n    return False\n\n  @property\n  def is_dynamic_size(self):\n    return any(not isinstance(i, int) for i in self.shape)\n\n  @functools.cached_property","sourceCodeStart":278,"sourceCodeEnd":314,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/types.py#L278-L314","documentation":"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.","triggerScenarios":"Calling chained methods like ref_group.transpose(...).reshape(...) or .bitcast(...) then another transform on a multiref.","commonSituations":"Trying to apply .transpose to a multiref, or bitcast+reshape chains on grouped refs.","solutions":["Apply the transform to each individual ref instead of the group","Flatten to a single ref before chaining transforms","Restructure so the multiref only ever receives one transform"],"exampleFix":"// before\ngroup.transpose(perm).reshape(new_shape)\n// after\nfor r in group_ref.ref:\n    r.transpose(perm).reshape(new_shape)","handlingStrategy":"validation","validationCode":"if is_multiref(ref) and ref.transforms:\n    raise ValueError(\"apply to individual refs instead\")","typeGuard":"def is_multiref(ref):\n    return isinstance(getattr(ref, \"ref\", None), tuple)","tryCatchPattern":null,"preventionTips":["Never chain transforms on multirefs","Apply per-ref transforms in loops"],"tags":["jax","multiref","api-restriction"],"backgroundTag":"unsupported-operation-chain","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}