jax-ml/jax · error · NotImplementedError

All aliased Refs must have the same memory space (SMEM or TM

Error message

All aliased Refs must have the same memory space (SMEM or TMEM). Got {(ref.memory_space for ref in ref_leaves)}.

What it means

RefMap in Mosaic GPU only supports aliasing refs that all live in the same memory space — either all SMEM or all TMEM. Mixing memory spaces (or including refs in other spaces) in one aliased union raises NotImplementedError. Note the message formats a generator expression, so the printed 'Got' value shows the generator repr, not the actual spaces.

Source

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

      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:
    inner_aval = jax.core.ShapedArray(self.shape, self.dtype)
    refs_aval = jax.tree.map(lambda ref: ref.get_ref_aval(), self.refs)
    return AbstractRefUnion(inner_aval, refs_aval,
                            memory_space=self.memory_space)


Index = mgpu.DynamicSlice | slice | int | ir.Value


@dataclasses.dataclass(frozen=True)
class TilingTransform(state_types.Transform):
  """Represents a tiling transformation for memory refs.

  A tiling of (X, Y) on an array of shape (M, N) will result in a transformed

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Split the aliased refs into separate RefMaps per memory space
If cross-space sharing was intended, remove it — copy explicitly between SMEM and TMEM instead of aliasing
Debug by printing `[ref.memory_space for ref in jax.tree.leaves(refs)]` before constructing the RefMap (the error message itself is unhelpful due to the generator formatting bug)

Example fix

# before
ref_map = RefMap((smem_ref, tmem_ref))  # NotImplementedError

# after
smem_map = RefMap((smem_ref, other_smem_ref))
tmem_map = RefMap((tmem_ref, other_tmem_ref))

Defensive patterns

Strategy: validation

Validate before calling

spaces = {r.memory_space for r in jax.tree.leaves(refs)}
assert len(spaces) == 1, f'mixed memory spaces: {spaces}'

Type guard

def refs_share_memory_space(refs) -> bool:
    spaces = {r.memory_space for r in jax.tree.leaves(refs)}
    return len(spaces) == 1

Try / catch

null

Prevention

Alias only within one memory space
Print [r.memory_space for r in leaves] when debugging — the error message hides them

When it happens

Trigger: Constructing a RefMap whose leaves contain e.g. one SMEM ref and one TMEM ref, or a ref in an unsupported space (e.g. HBM), so neither the all-SMEM nor the all-TMEM branch of __init__ applies.

Common situations: Aliasing scratch buffers across memory spaces in Pallas kernels (e.g. sharing an SMEM staging buffer with a TMEM accumulator); refactors of buffer reuse/aliasing passes; passing a mixed pytree where a RefMap accidentally includes an extra ref.

Related errors

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