jax-ml/jax · error · NotImplementedError

Replicated dimensions are not supported

Error message

Replicated dimensions are not supported

What it means

store_tiled_async emits per-warp/lane PTX that assumes each register dimension maps to exactly one thread; a Replicated dimension means multiple threads hold the same element, and the store path has no vote/serialization to handle that, so it is rejected.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3806

      cluster_dim: gpu.Dimension,
      cluster_idx: ir.Value,
      swizzle: int | None,
      optimized: bool = True,
      tiling_rank: int | None = None,
      atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
  ):
    i32 = ir.IntegerType.get_signless(32)
    i64 = ir.IntegerType.get_signless(64)
    if isinstance(ref, utils.MultimemRef):
      raise ValueError("Multimem refs are not supported in store_tiled_async")
    layout, shape = self.layout, self.shape
    if not isinstance(layout, TiledLayout):
      raise NotImplementedError(self.layout)
    if any(
        isinstance(d, Replicated)
        for d in itertools.chain(layout.warp_dims, layout.lane_dims)
    ):
      raise NotImplementedError("Replicated dimensions are not supported")
    full_cluster_idx: list[ir.Value] = [
        gpu.cluster_block_id(d) for d in gpu.Dimension
    ]
    full_cluster_idx[cluster_dim] = cluster_idx
    lin_cluster_idx = arith.index_cast(
        i32, utils.cluster_idx(tuple(gpu.Dimension), full_cluster_idx)
    )
    cluster_barrier_ptr = utils.get_cluster_ptr(
        barrier.get_ptr(), lin_cluster_idx, generic=False
    )
    cluster_ref = utils.get_cluster_ref(
        ref, cluster_dim, cluster_idx, generic=False
    )
    stores = self.transfer_tiled(
        cluster_ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank
    )
    if atomic is not None:
      for get, _update, _idx, cluster_ptr in stores:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a layout without replicated warp/lane dims before storing: fa.to_layout(TiledLayout(vec_size=..., replicates=None))
2. Re-fragment the value so each element lives in exactly one thread
3. Fall back to store_tiled / non-async stores for replicated layouts

Example fix

// before
fa.store_tiled_async(ref, swizzle=1)
// after
fa = fa.to_layout(fa.layout.to_tiled())  # drop Replicated dims
fa.store_tiled_async(ref, swizzle=1)

import itertools
from jax.experimental.mosaic.gpu.fragmented_array import Replicated
if any(isinstance(d, Replicated) for d in itertools.chain(fa.layout.warp_dims, fa.layout.lane_dims)):
    fa = fa.to_layout(fa.layout.to_tiled())
fa.store_tiled_async(ref, ...)

from jax.experimental.mosaic.gpu.fragmented_array import Replicated
import itertools

def is_async_storeable(fa) -> bool:
    l = fa.layout
    return not any(isinstance(d, Replicated) for d in itertools.chain(l.warp_dims, l.lane_dims))

try:
    fa.store_tiled_async(ref, ...)
except NotImplementedError as e:
    if 'Replicated' in str(e):
        fa.to_layout(fa.layout.to_tiled()).store_tiled_async(ref, ...)
    else:
        raise

When it happens

Trigger: Calling store_tiled_async when any of layout.warp_dims or layout.lane_dims is a Replicated instance (layouts built with Replicated in the lane/warp mapping, or layouts emerging from reductions that keep a replicated dim).

Common situations: Using a layout with replicated lanes after ops that produce replicated results; specifying a layout with Replicated(..., dim=...) in warp/lane dims when constructing arrays for async stores.

Related errors

• Loading multiple row tiles
• {op} has an unsupported layout: {out_layout_attr}
• Expected TiledLayout, got {type(layout)}
• Output layout {out_layout} must match the accumulator layout
• Unsupported layout: {src.layout}

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