jax-ml/jax · error · TypeError

ref must be a reference

Error message

ref must be a reference

What it means

remote_ref requires its argument to be a reference (an AbstractRef or TransformedRef). Passing a plain array, tracer, or other value raises TypeError('ref must be a reference').

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1053

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    raise NotImplementedError()

  def commute_ndindexer(
      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, MulticastRef]:
    return indexer, self


def remote_ref(
    ref: _Ref,
    device_id: jax.typing.ArrayLike,
    device_id_type: pallas_primitives.DeviceIdType = pallas_primitives.DeviceIdType.MESH,
) -> pallas_core.TransformedRef:
  """Translate memref to a symmetric memref on a peer device."""
  if not isinstance(ref, pallas_core.TransformedRef):
    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):
      raise TypeError("ref must be a reference")
    ref = pallas_core.TransformedRef(ref, transforms=())
  if any(isinstance(t, MulticastRef) for t in ref.transforms):
    raise ValueError("Can't make a multicast reference into a peer reference.")
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, PeerMemRef(device_id, device_id_type)),
  )


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ClusterRefTransform(state_types.Transform):
  dims: tuple[jax_core.AxisName, ...] = jax.tree.static()
  idxs: tuple[Any, ...]

  def __post_init__(self):
    if len(self.dims) != len(self.idxs):
      raise ValueError("dims and idxs must have the same length")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a pallas Ref / TransformedRef obtained from the kernel's input/output parameters
2. Ensure the value comes from the kernel signature (Ref[...]) rather than a computed array
3. Wrap raw refs consistently: let remote_ref wrap plain AbstractRefs itself

Example fix

// before
peer = remote_ref(some_array, peer_id)

// after
peer = remote_ref(out_ref, peer_id)  # out_ref is a Ref from kernel params

from jax._src.pallas import state_types
assert isinstance(ref, pallas_core.TransformedRef) or isinstance(jax_core.typeof(ref), state_types.AbstractRef)

def is_ref(x): return isinstance(x, pallas_core.TransformedRef) or type(jax_core.typeof(x)).__name__ == 'AbstractRef'

When it happens

Trigger: Calling pallas_mosaic_gpu.remote_ref(array, device_id) with a non-ref value, e.g. passing a block from a kernel argument that isn't a Ref, or passing an already-consumed value.

Common situations: Writing cross-device TPU kernels and passing jax arrays or scalars where a memref is expected; passing a ref created by a different pallas backend.

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

• Can't make a multicast reference into a peer reference.
• Can't make a peer reference into a multicast reference.
• Unsupported type: {x}
• Can only store to references (got {x_ref}).
• Can only store scalars or vectors (got {value}).

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