jax-ml/jax · error · ValueError

Unsupported memory space: {orig_ref_ty.memory_space}

Error message

Unsupported memory space: {orig_ref_ty.memory_space}

What it means

The vector load lowering hit a reference whose memory space is neither SMEM nor plain GMEM-with-transforms, i.e. transforms_attr is None for a memory space it cannot handle. Only GMEM with an attached transforms attribute (tiling/swizzle) is supported on this path.

Source

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

        vec_size=out_layout.vec_size,
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if not isinstance(out_layout, fa.TiledLayout):
    raise ValueError(f"{op} has an unsupported layout: {out_layout_attr}")

  optimized = op.optimized.value if op.optimized is not None else None
  if transformed_ref.type.memory_space is None:  # GMEM
    fragmented_array = fa.FragmentedArray.load_untiled(
        transformed_ref,
        layout=out_layout,
        is_signed=is_signed,
        optimized=bool(optimized),
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if transforms_attr is None:
    raise ValueError(f"Unsupported memory space: {orig_ref_ty.memory_space}")

  swizzle = swizzle_from_transforms_attr(transforms_attr)
  transforms = memref_transforms_from_transforms_attr(transforms_attr)
  if transforms:
    [tiling_transform] = transforms
    assert isinstance(tiling_transform, lc.TileTransform)

    def load_tiled(optimized: bool) -> fa.FragmentedArray:
      return fa.FragmentedArray.load_tiled(
          transformed_ref,
          swizzle,
          is_signed=is_signed,
          layout=out_layout,
          optimized=optimized,
          tiling_rank=len(tiling_transform.tiling)
      )

    fragmented_array = _retry_on_failure(load_tiled, optimized)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Load from shared memory or a plain GMEM ref created by mosaic's alloc helpers
2. Ensure the ref carries in_transforms (tiling/swizzle) if loading from GMEM
3. Avoid mgpu.memref_load style ops on exotic memory spaces; use mosaic's smem utilities

space = ref.type.memory_space
assert space is None or int(space) == 1 or transforms_attr is not None, f'unsupported memory space {space}'

def loadable_ref(ref, transforms_attr) -> bool:
    return ref.type.memory_space is None and transforms_attr is not None

When it happens

Trigger: Loading from a memref whose memory_space is set (e.g. #mgpu.thread or another address space) or a GMEM ref without a transforms attribute, hitting _vector_load_op_lowering_rule's fallback branch.

Common situations: Loading directly from a reference to special address spaces (e.g. tensor-memory or thread-local memory) instead of SMEM; constructing refs manually via memref ops bypassing Mosaic helpers.

Related errors

• Unsupported memory space: {ref_type.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/39711c3eb35228b6. Report an issue: GitHub.