jax-ml/jax · error · ValueError

Unsupported memory space: {ref_type.memory_space}

Error message

Unsupported memory space: {ref_type.memory_space}

What it means

The vector store lowering only supports SMEM (memory_space == 1 path) and the GMEM/transforms path; any other memory space on the destination ref raises this ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:583

      assert isinstance(tiling_transform, lc.TileTransform)

      def store_tiled(optimized: bool):
        fragmented_array.store_tiled(
            unwrapped_ref, swizzle, optimized,
            tiling_rank=len(tiling_transform.tiling),
            atomic=atomic
        )

      _retry_on_failure(store_tiled, optimized)
    else:

      def store_untiled(optimized: bool):
        fragmented_array.store_untiled(
            ref, optimized=optimized, swizzle=swizzle, atomic=atomic
        )
      _retry_on_failure(store_untiled, optimized)
  else:
    raise ValueError(f"Unsupported memory space: {ref_type.memory_space}")

  if ctx.auto_barriers:
    utils.warpgroup_barrier()  # Make sure the writes have completed.

  return []


@_register_lowering(mgpu.AsyncStoreSmemOp)
def _async_store_smem_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.AsyncStoreSmemOp
) -> Sequence[ir.Value]:
  index = ir.IndexType.get()

  [to_store_layout] = inference_utils.in_layouts(op)
  value = _fragmented_array_from_ir(op.valueToStore, to_store_layout)
  layout = layouts_lib.from_layout_attr(to_store_layout)
  if not isinstance(layout, fa.TiledLayout):
    raise NotImplementedError(f"Expected TiledLayout, got {type(layout)}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Store to shared memory or a plain global-memory ref with proper transforms
2. Allocate destinations with mosaic's alloc/smem helpers rather than raw memref.alloc with custom memory spaces

space = ref.type.memory_space
assert space is None or int(space) == 1, f'cannot store to memory space {space}'

def storable_ref(ref) -> bool:
    return ref.type.memory_space is None or int(ref.type.memory_space) == 1

When it happens

Trigger: Storing a fragmented array to a memref with memory_space other than SMEM or None/GMEM-with-transforms in _vector_store_op_lowering_rule.

Common situations: Storing to thread-local or tensor-memory refs; using custom memref allocations with unusual address space attributes in a Mosaic kernel.

Related errors

• Unsupported memory space: {orig_ref_ty.memory_space}
• Memory space {self.memory_space} is not supported by mesh {s
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Expected TiledLayout, got {type(layout)}

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