jax-ml/jax · error · ValueError

Some aliased TMEM references are collective and some are not

Error message

Some aliased TMEM references are collective and some are not.

What it means

RefMap (a union of aliased refs) in Mosaic GPU requires that when all aliased refs live in TMEM, they must all have the same `collective` flag, since the union exposes a single collective value. Mixing collective and non-collective TMEM refs is ambiguous and raises ValueError at construction.

Source

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

  def __init__(self, *refs: _GPUMemoryRefTree):
    ref_leaves = jax.tree.leaves(refs)
    if all(ref.memory_space == SMEM for ref in ref_leaves):
      object.__setattr__(self, "refs", refs)
      num_bytes = max(map(_ref_group_size, self.refs))
      super().__init__(
          inner_aval=jax_core.ShapedArray(
              (num_bytes,), jnp.int8
          ),
          memory_space=SMEM,
          transforms=(),
      )
    elif all(ref.memory_space == TMEM for ref in ref_leaves):
      object.__setattr__(self, "refs", refs)
      max_cols = max(map(_ref_group_tmem_col_size, self.refs))
      is_collective = ref_leaves[0].collective
      if any(r.collective != is_collective for r in ref_leaves):
        raise ValueError(
            "Some aliased TMEM references are collective and some are not."
        )
      super().__init__(
          inner_aval=jax_core.ShapedArray(
              shape=(128, max_cols,),
              dtype=jnp.int32,
          ),
          memory_space=TMEM,
          transforms=(),
          layout=tcgen05.tmem_default_layout(packing=1),
          collective=all(ref.collective for ref in ref_leaves),
      )
    else:
      raise NotImplementedError(
          "All aliased Refs must have the same memory space (SMEM or TMEM). "
          f"Got {(ref.memory_space for ref in ref_leaves)}.")

  def get_ref_aval(self) -> AbstractRefUnion:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create all TMEM refs that will be aliased with the same collective mode (pass matching collective= argument to the tcgen05 allocation)
2. Split the RefMap so collective and non-collective refs are not aliased together
3. Check `ref.collective` on each TMEM ref before building the RefMap and adjust allocation

Example fix

# before
ref_a = tcgen05.alloc(..., collective=True)
ref_b = tcgen05.alloc(..., collective=False)
ref_map = RefMap((ref_a, ref_b))  # ValueError

# after
ref_b = tcgen05.alloc(..., collective=True)
ref_map = RefMap((ref_a, ref_b))

leaves = jax.tree.leaves(refs)
assert all(
    r.memory_space != mosaic_gpu_core.TMEM
    or r.collective == leaves[0].collective
    for r in leaves
), 'mixed collective modes in aliased TMEM refs'

def refs_have_uniform_collective(refs) -> bool:
    leaves = jax.tree.leaves(refs)
    tmem = [r for r in leaves if r.memory_space == mosaic_gpu_core.TMEM]
    return len({r.collective for r in tmem}) <= 1

null

When it happens

Trigger: Constructing a RefMap (RefUnion) whose pytree leaves include TMEM refs created with different collective settings, e.g. combining refs allocated via `tcgen05.alloc` collective API with per-thread TMEM refs.

Common situations: Refactoring kernels to share/alias TMEM buffers between collective matmul ops and per-warpgroup ops; version upgrades that introduced the collective invariant; tests constructing mixed ref trees for aliasing analysis.

Related errors

• All aliased Refs must have the same memory space (SMEM or TM
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• layout attribute is only defined for TMEM refs
• collective attribute is only defined for TMEM refs

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