{"record":{"id":"2600188fa111cd02","repo":"jax-ml/jax","slug":"cannot-eagerly-run-with-memory-space-constraint","errorCode":null,"errorMessage":"Cannot eagerly run with_memory_space_constraint.","messagePattern":"Cannot eagerly run with_memory_space_constraint\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":1618,"sourceCode":"    raise ValueError(f\"Mesh {self=} is not compatible with {other_mesh=}.\")\n\n  @property\n  def supported_memory_spaces(self) -> Sequence[Any]:\n    \"\"\"Return the memory spaces supported by the mesh.\"\"\"\n\n  @contextlib.contextmanager\n  def tracing_context(self) -> Generator[None]:\n    raise NotImplementedError()\n    yield\n\n\nwith_memory_space_constraint_p = jax_core.Primitive(\n    'with_memory_space_constraint')\n\n@with_memory_space_constraint_p.def_impl\ndef with_memory_space_constraint_impl(x, *, memory_space):\n  del x, memory_space\n  raise ValueError(\"Cannot eagerly run with_memory_space_constraint.\")\n\n\n@with_memory_space_constraint_p.def_abstract_eval\ndef with_memory_space_constraint_abstract_eval(x, *, memory_space):\n  if not isinstance(x, jax_core.ShapedArray):\n    raise NotImplementedError(\"with_memory_space_constraint only supports \"\n                              \"arrays.\")\n  return x.update(memory_space=memory_space)\n\ndef with_memory_space_constraint_lowering_rule(ctx, x, *, memory_space):\n  del ctx, memory_space\n  return [x]\nmlir.register_lowering(\n    with_memory_space_constraint_p, with_memory_space_constraint_lowering_rule\n)\n\n\ndef with_memory_space_constraint_batching_rule(","sourceCodeStart":1600,"sourceCodeEnd":1636,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L1600-L1636","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\ny = with_memory_space_constraint(x, memory_space=MemorySpace.VMEM)  # eager -> ValueError\n\n# after\ny = jax.jit(lambda x: with_memory_space_constraint(x, memory_space=MemorySpace.VMEM))(x)","handlingStrategy":"validation","validationCode":"import jax\n# ensure tracing: only call under jit\nassert jax.config.read('jax_disable_jit') is False or _under_trace(), 'call inside jit/pallas only'","typeGuard":"null","tryCatchPattern":"try:\n    result = with_memory_space_constraint(x, memory_space=ms)\nexcept ValueError as e:\n    if 'Cannot eagerly run' in str(e):\n        result = jax.jit(lambda v: with_memory_space_constraint(v, memory_space=ms))(x)\n    else:\n        raise","preventionTips":["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"],"tags":["jax","pallas","eager-execution","primitive","memory-space"],"backgroundTag":"jax-primitive-eager-execution","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}