jax-ml/jax · error · TypeError

Cannot select from Refs of different types: {types}

Error message

Cannot select from Refs of different types: {types}

What it means

SelectTransform.transform_types raises TypeError when the refs being selected from have differing abstract types, since a single selection result type can't be chosen. All refs in a multiref select must share the same aval.

Source

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

@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)
class SelectTransform(MultiRefTransform):
  idx: Array | int

  def transform_types(self, xs):
    def _type(ref):
      match ref:
        case AbstractRef():
          return ref
        case core.ShapedArray():
          raise NotImplementedError
        case _:
          raise TypeError(f"Cannot select {ref}")

    assert isinstance(xs, Sequence), f"Select expected sequence, got {xs}"
    types = tuple(_type(ref) for ref in xs)
    if any(types[0] != t for t in types[1:]):
      raise TypeError(f"Cannot select from Refs of different types: {types}")
    return types[0]

  def undo(self, x: core.AbstractValue) -> Transform:
    raise NotImplementedError(type(self))

  def pretty_print(self, context: core.JaxprPpContext) -> pp.Doc:
    del context  # Unused.
    return pp.text(f"{{select({self.idx=})}}")

  def getattr(self, name: str, xs: Sequence[core.AbstractValue]) -> Any:
    attrs = [getattr(x, name) for x in xs]
    if any(attrs[0] != attr for attr in attrs[1:]):
      raise TypeError(f"Cannot resolve attribute {name} from: {attrs}")
    return attrs[0]


@dataclasses.dataclass(slots=True)
class RefIndexer:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give all refs in the group the same dtype and shape
2. Split the select into per-type groups
3. Assert type homogeneity before building the multiref

Example fix

// before
refs = (f32_ref, i32_ref)
refs.get()  # TypeError
// after
refs = (f32_ref, f32_ref)
refs.get()

types = {core.typeof(r) for r in refs}
assert len(types) == 1, f"heterogeneous refs: {types}"

try:
    refs.get()
except TypeError as e:
    if "different types" in str(e):
        # split group by type
        ...

When it happens

Trigger: Calling get/put on a multiref whose constituent refs have different dtypes or shapes, e.g. mixing a float32 ref and an int32 ref in one TransformedRef.

Common situations: Grouping heterogeneous buffers into one select for convenience; dtype changes in one buffer but not others after a refactor.

Related errors

• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• input type mismatch for {_prim}
• at {keystr(path)}, got fwd output type {ty.str_short()} whic
• {result}the bwd rule attached to {primal_sourceinfo} produce
• mulhi operands must have the same dtype, got {dtype} and {y_

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