jax-ml/jax · error · NotImplementedError

Unsupported memory space.

Error message

Unsupported memory space.

What it means

Terminal fallback in _extract_aliased_ref: the aliased ref's memory space is neither SMEM (smem()) nor TMEM (tmem()), so no aliasing strategy exists. The lowering switches on the ref's memory space and raises NotImplementedError for anything else, such as global/device memory or WMEM.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:1670

          total_offset = base_offset + offset
          ref_ty = ir.MemRefType.get(
              transformed_shape, mlir_dtype, memory_space=mgpu_utils.tmem()
          )
          alloc_id = source_slice_op.alias_id
          assert alloc_id is not None
          # TODO(bchetioui): Use a scheme resilient to hash collisions.
          alias_id = hash((offset, alloc_id.value, alias_group_idx))
          slice_op = mgpu.dialect.SliceTmemOp(
              ref_ty, source_slice_op.source, total_offset
          )
          i64 = ir.IntegerType.get_signless(64)
          slice_op.attributes["alias_id"] = ir.IntegerAttr.get(i64, alias_id)
          ref = slice_op.result
          assert layout is not None
          layout_attr = mgpu.layouts.to_layout_attr(layout)
          ref = mgpu.dialect.tmem_layout_cast(ref, layout_attr)
        else:
          raise NotImplementedError("Unsupported memory space.")
      return (
          ref,
          ref_aval,
          transform_avals[1:],
          tuple(other_transforms),
      )
    case _:
      # No ExtractAliasedRef found, don't do anything.
      return ref, ref_aval, transform_avals, transforms


def _commute_transform(
    aval: jax_core.AbstractValue,
    t1: state_types.Transform,
    t2: state_types.Transform,
) -> tuple[state_types.Transform, state_types.Transform]:
  """Commutes two transforms.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Copy the global/wmem data into SMEM scratch first, then alias/bitcast the SMEM buffer
2. Remove view/aliasing from refs that live in unsupported memory spaces; convert dtypes on values instead
3. Check that the ref you're viewing was actually allocated by SMEM/TMEM helpers and not passed in from global memory

Example fix

# before
scratch = in_ref.view(jnp.float32)  # in_ref is global memory -> error
# after
smem = alloc_smem(in_ref.shape, jnp.uint8)
smem[...] = in_ref.astype(jnp.uint8)
scratch = smem.view(jnp.float32)

space = getattr(buf, 'memory_space', None)
assert space in ('smem', 'tmem'), f'cannot alias ref in memory space {space!r}'

def aliasable_memory_space(buf) -> bool:
    return getattr(buf, 'memory_space', None) in ('smem', 'tmem')

When it happens

Trigger: Creating an aliased/viewed Ref that resolves to a memory space other than shared or tensor memory — e.g. aliasing a global-memory buffer or a warpgroup-register (wmem) ref inside a Pallas GPU kernel.

Common situations: Kernel code that aliases input/output refs directly (which live in global memory) instead of staging through SMEM; new memory spaces introduced by newer Mosaic dialects.

Related errors

• Only SMEM and TMEM refs are supported.
• The base ref for aliases must come from a slice_smem op.
• Unsupported memory space: {x.memory_space}
• Unsupported dtype: {ref.dtype}
• 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/ab97a0aa236ebe1d. Report an issue: GitHub.