jax-ml/jax · error · TypeError

Cannot select {ref}

Error message

Cannot select {ref}

What it means

Raised by SelectTransform._type when an element of the sequence being selected from is neither an AbstractRef nor a ShapedArray. JAX raises TypeError because it cannot compute a common type for the selection.

Source

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

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


@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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify every element of the ref tuple is a ref (has .aval of AbstractRef)
2. Filter out non-ref values before constructing the multiref
3. Check for accidental nesting, e.g. a tuple inside the tuple

assert all(type(r).__name__ in ("AbstractRef",) or hasattr(r, 'inner_aval') for r in refs), "all elements must be refs"

def all_refs(xs):
    return all(hasattr(x, "inner_aval") or hasattr(x, "aval") for x in xs)

When it happens

Trigger: Passing a tuple/list containing tokens or other non-ref, non-array avals to a select transform over multiple refs.

Common situations: Building multirefs with heterogeneous elements; custom jaxpr plumbing that mixes ref and non-ref values.

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

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

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