jax-ml/jax · error · ValueError

Can't transpose a TMEM reference.

Error message

Can't transpose a TMEM reference.

What it means

transpose_ref rejects transposing references whose memory_space is TMEM (tensor memory). TMEM layout on TPU cannot be permuted as an address transform, so the operation raises ValueError.

Source

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

    collective_axes = (collective_axes,)
  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, PeerMemRef) for t in ref.transforms):
    raise ValueError("Can't make a peer reference into a multicast reference.")
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, MulticastRef(collective_axes)),
  )


def transpose_ref(
    ref: pallas_core.TransformedRef | Any,
    permutation: tuple[int, ...],
) -> pallas_core.TransformedRef:
  assert hasattr(ref, "memory_space")
  if ref.memory_space == MemorySpace.TMEM:
    raise ValueError("Can't transpose a TMEM reference.")
  return ref.transpose(permutation)


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ExtractAliasedRef(state_types.Transform):
  """Bitcasts the underlying ref at the given offset to the given shape and dtype."""
  dtype: dtypes.DType = jax.tree.static()
  shape: tuple[int, ...] = jax.tree.static()
  offset: int = jax.tree.static()

  # The index of the group of this aliased ref within the input RefUnion.
  alias_group_idx: int = jax.tree.static()

  # TMEM-specific params
  layout: tcgen05.TMEMLayout | None = jax.tree.static()

  @classmethod

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Skip transposition for TMEM refs: check ref.memory_space first
2. Use explicit TMEM-aware layout ops (e.g. load/store with transposed indices) instead of transpose_ref
3. Move data out of TMEM into SMEM/VMEM before transposing

Example fix

# before
ref = transpose_ref(tmem_ref, (1, 0))

# after
if ref.memory_space != MemorySpace.TMEM:
    ref = transpose_ref(ref, (1, 0))

from jax._src.pallas.mosaic_gpu.core import MemorySpace
if ref.memory_space == MemorySpace.TMEM: raise ValueError('cannot transpose TMEM')

def is_transposable(ref): return getattr(ref, 'memory_space', None) != MemorySpace.TMEM

When it happens

Trigger: Calling transpose_ref on a ref allocated in MemorySpace.TMEM (e.g. a tmem_t allocation or TMEM block) with any permutation.

Common situations: Kernels using tensor-core TMEM buffers (tcgen05/tmem) where generic transpose helper code is applied to all refs indiscriminately.

Related errors

• 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
• Some aliased TMEM references are collective and some are not

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