jax-ml/jax · error · NotImplementedError

Unsupported memory space: {x.memory_space}

Error message

Unsupported memory space: {x.memory_space}

What it means

The remove_memory_space abstract eval strips memory-space annotations from arrays, but only supports memory_space of None, ANY, or HBM (mosaic). Any other memory space annotation raises NotImplementedError, since the interpreter cannot erase an unrecognized memory space.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1896

          subkey, coordindates_along_dim
      )
      grid_point_coordinates.append(coordindates_along_dim)
    else:
      grid_point_coordinates.append(jnp.array((), dtype=jnp.int32))

  return tuple(grid_point_coordinates)

# TODO(sharadmv, jburnim): add support for memory space constraints
remove_memory_space_p = jax_core.Primitive('remove_memory_space')

@remove_memory_space_p.def_abstract_eval
def _remove_memory_space_abstract_eval(x):
  if (isinstance(x, jax_core.ShapedArray) and
      not isinstance(x.memory_space, jax_core.MemorySpace)):
    if (x.memory_space is None or x.memory_space is pallas_core.MemorySpace.ANY
        or x.memory_space is mosaic_core.MemorySpace.HBM):
      return jax_core.ShapedArray(x.shape, x.dtype)
    raise NotImplementedError(f'Unsupported memory space: {x.memory_space}')
  return x

@remove_memory_space_p.def_impl
def _remove_memory_space_impl(x):
  return x

def _remove_memory_space_lowering(_, x):
  return [x]
mlir.register_lowering(remove_memory_space_p, _remove_memory_space_lowering)


def _get_grid_point(
    loop_indices: tuple[Array, ...],
    grid_point_coordinates: _GridPointCoordinatesPerDim,
) -> Array:
  """Indexes each entry in `grid_point_coordinates` with the corresponding entry in `loop_indices`.

  If an entry in `grid_point_coordinates` is an empty array, the corresponding

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass plain jax arrays (no memory space annotation) as pallas_call inputs/outputs
2. Allocate VMEM/SMEM inside the kernel via run_scoped, not as external arguments
3. Update JAX/mosaic versions so supported memory spaces agree

Example fix

# before
outs = pallas_call(kernel, out_shapes, vmem_annotated_inputs)  # annotated arrays
# after
outs = pallas_call(kernel, out_shapes, plain_jnp_arrays)

for a in jax.tree_util.tree_leaves(args):
    ms = getattr(a, 'memory_space', None)
    assert ms is None or 'ANY' in str(ms) or 'HBM' in str(ms), f'unsupported memory_space {ms}'

def has_supported_memory_space(a) -> bool:
    ms = getattr(a, 'memory_space', None)
    return ms is None or 'ANY' in str(ms) or 'HBM' in str(ms)

try:
    interpret_run(kernel, args)
except NotImplementedError as e:
    if 'Unsupported memory space' in str(e):
        args = jax.tree.map(lambda x: jnp.asarray(x) if hasattr(x, 'memory_space') else x, args)
        interpret_run(kernel, args)

When it happens

Trigger: Passing an array whose memory_space attribute is an unsupported enum value (e.g., a mosaic VMEM/SMEM annotated value) through the remove_memory_space_p primitive during interpretation setup.

Common situations: Feeding pre-annotated buffers (VMEM refs) directly as pallas_call inputs/outputs instead of plain HBM arrays; version skew where new memory-space enums exist in mosaic but not the interpreter path.

Related errors

• Only SMEM and TMEM refs are supported.
• Unsupported memory space.
• array ref with memory space only works inside of a `jit`.
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton

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