jax-ml/jax · error · ValueError

Cannot eagerly run with_memory_space_constraint.

Error message

Cannot eagerly run with_memory_space_constraint.

What it means

with_memory_space_constraint is a JAX primitive used inside Pallas/MLIR lowering pipelines to tag a value with a target memory space (e.g. VMEM/SMEM). It has no eager implementation by design; calling it directly on concrete values raises ValueError. It only works under tracing (jit/pallas kernel tracing) where lowering rules consume it.

Source

Thrown at jax/_src/pallas/core.py:1618

    raise ValueError(f"Mesh {self=} is not compatible with {other_mesh=}.")

  @property
  def supported_memory_spaces(self) -> Sequence[Any]:
    """Return the memory spaces supported by the mesh."""

  @contextlib.contextmanager
  def tracing_context(self) -> Generator[None]:
    raise NotImplementedError()
    yield


with_memory_space_constraint_p = jax_core.Primitive(
    'with_memory_space_constraint')

@with_memory_space_constraint_p.def_impl
def with_memory_space_constraint_impl(x, *, memory_space):
  del x, memory_space
  raise ValueError("Cannot eagerly run with_memory_space_constraint.")


@with_memory_space_constraint_p.def_abstract_eval
def with_memory_space_constraint_abstract_eval(x, *, memory_space):
  if not isinstance(x, jax_core.ShapedArray):
    raise NotImplementedError("with_memory_space_constraint only supports "
                              "arrays.")
  return x.update(memory_space=memory_space)

def with_memory_space_constraint_lowering_rule(ctx, x, *, memory_space):
  del ctx, memory_space
  return [x]
mlir.register_lowering(
    with_memory_space_constraint_p, with_memory_space_constraint_lowering_rule
)


def with_memory_space_constraint_batching_rule(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Wrap the code path in jax.jit (or run it inside a pallas kernel trace) so the primitive is staged and lowered instead of executed
Avoid calling the private jax._src.pallas.core.with_memory_space_constraint directly; use the public pallas API for memory-space hints
If you need an eager no-op for testing, guard with jax.core.eval_context or substitute an identity function in tests

Example fix

# before
y = with_memory_space_constraint(x, memory_space=MemorySpace.VMEM)  # eager -> ValueError

# after
y = jax.jit(lambda x: with_memory_space_constraint(x, memory_space=MemorySpace.VMEM))(x)

Defensive patterns

Strategy: validation

Validate before calling

import jax
# ensure tracing: only call under jit
assert jax.config.read('jax_disable_jit') is False or _under_trace(), 'call inside jit/pallas only'

Type guard

null

Try / catch

try:
    result = with_memory_space_constraint(x, memory_space=ms)
except ValueError as e:
    if 'Cannot eagerly run' in str(e):
        result = jax.jit(lambda v: with_memory_space_constraint(v, memory_space=ms))(x)
    else:
        raise

Prevention

Never call private _src primitives eagerly; always stage them under jax.jit or a pallas kernel
Keep memory-space constraint calls inside kernel body functions, not host code
In tests, substitute identity functions for primitives that lack eager impls

When it happens

Trigger: Directly calling jax._src.pallas.core.with_memory_space_constraint (or a helper that uses it) on a concrete ndarray outside of a traced/jitted context; printing or eagerly evaluating a function containing this primitive.

Common situations: Refactoring Pallas kernel helpers so they run eagerly during debugging; calling a library function that internally uses the primitive without wrapping the call in jax.jit or a pallas kernel; using the private _src API instead of a public wrapper.

Related errors

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