{"record":{"id":"802a7e4900afaf79","repo":"jax-ml/jax","slug":"replicated-dimensions-are-not-supported","errorCode":null,"errorMessage":"Replicated dimensions are not supported","messagePattern":"Replicated dimensions are not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3806,"sourceCode":"      cluster_dim: gpu.Dimension,\n      cluster_idx: ir.Value,\n      swizzle: int | None,\n      optimized: bool = True,\n      tiling_rank: int | None = None,\n      atomic: Literal[\"add\", \"max\", \"min\", \"and\", \"or\", \"xor\"] | None = None,\n  ):\n    i32 = ir.IntegerType.get_signless(32)\n    i64 = ir.IntegerType.get_signless(64)\n    if isinstance(ref, utils.MultimemRef):\n      raise ValueError(\"Multimem refs are not supported in store_tiled_async\")\n    layout, shape = self.layout, self.shape\n    if not isinstance(layout, TiledLayout):\n      raise NotImplementedError(self.layout)\n    if any(\n        isinstance(d, Replicated)\n        for d in itertools.chain(layout.warp_dims, layout.lane_dims)\n    ):\n      raise NotImplementedError(\"Replicated dimensions are not supported\")\n    full_cluster_idx: list[ir.Value] = [\n        gpu.cluster_block_id(d) for d in gpu.Dimension\n    ]\n    full_cluster_idx[cluster_dim] = cluster_idx\n    lin_cluster_idx = arith.index_cast(\n        i32, utils.cluster_idx(tuple(gpu.Dimension), full_cluster_idx)\n    )\n    cluster_barrier_ptr = utils.get_cluster_ptr(\n        barrier.get_ptr(), lin_cluster_idx, generic=False\n    )\n    cluster_ref = utils.get_cluster_ref(\n        ref, cluster_dim, cluster_idx, generic=False\n    )\n    stores = self.transfer_tiled(\n        cluster_ref, swizzle, layout, shape, optimized, ref_tiling_rank=tiling_rank\n    )\n    if atomic is not None:\n      for get, _update, _idx, cluster_ptr in stores:","sourceCodeStart":3788,"sourceCodeEnd":3824,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3788-L3824","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Convert to a layout without replicated warp/lane dims before storing: fa.to_layout(TiledLayout(vec_size=..., replicates=None))","Re-fragment the value so each element lives in exactly one thread","Fall back to store_tiled / non-async stores for replicated layouts"],"exampleFix":"// before\nfa.store_tiled_async(ref, swizzle=1)\n// after\nfa = fa.to_layout(fa.layout.to_tiled())  # drop Replicated dims\nfa.store_tiled_async(ref, swizzle=1)","handlingStrategy":"validation","validationCode":"import itertools\nfrom jax.experimental.mosaic.gpu.fragmented_array import Replicated\nif any(isinstance(d, Replicated) for d in itertools.chain(fa.layout.warp_dims, fa.layout.lane_dims)):\n    fa = fa.to_layout(fa.layout.to_tiled())\nfa.store_tiled_async(ref, ...)","typeGuard":"from jax.experimental.mosaic.gpu.fragmented_array import Replicated\nimport itertools\n\ndef is_async_storeable(fa) -> bool:\n    l = fa.layout\n    return not any(isinstance(d, Replicated) for d in itertools.chain(l.warp_dims, l.lane_dims))","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, ...)\nexcept NotImplementedError as e:\n    if 'Replicated' in str(e):\n        fa.to_layout(fa.layout.to_tiled()).store_tiled_async(ref, ...)\n    else:\n        raise","preventionTips":["Avoid Replicated dims in layouts destined for async stores","Check warp/lane dims after reductions that may replicate"],"tags":["jax","mosaic","gpu","layout","replicated","warp","not-implemented"],"backgroundTag":"unsupported-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}