jax-ml/jax · error · NotImplementedError

with_memory_space_constraint only supports arrays.

Error message

with_memory_space_constraint only supports arrays.

What it means

During abstract evaluation, with_memory_space_constraint requires its input to be a ShapedArray so it can attach a memory_space annotation via aval.update. Passing a token, other abstract value type, or a non-array tracer triggers NotImplementedError. The memory-space annotation machinery simply does not support non-array abstract values.

Source

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

  @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(
    axis_data, batched_args, batch_dims, *, memory_space
):
  del axis_data  # Unused; the constraint does not depend on the mapped axis.
  (x,), (bdim,) = batched_args, batch_dims
  out = with_memory_space_constraint_p.bind(x, memory_space=memory_space)
  return out, bdim  # the computed value and where the batch axis ended up in it

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only apply with_memory_space_constraint to array-valued tracers; route tokens/other values around it
2. Check the intermediate value with jax.core.get_aval(x) to confirm it is a ShapedArray before applying the constraint
3. Upgrade jax — support for additional aval types may have been added in newer releases
4. Reorder your computation so the constraint is applied before any transformation that changes the aval type

Example fix

# before
out = with_memory_space_constraint(token_or_exotic, memory_space=ms)

# after
from jax.core import ShapedArray
if isinstance(jax.core.get_aval(v), ShapedArray):
    out = with_memory_space_constraint(v, memory_space=ms)
else:
    out = v  # leave non-array values unconstrained

import jax.core as jc
aval = jc.get_aval(x)
assert isinstance(aval, jc.ShapedArray), f'unsupported aval {type(aval).__name__}'

def is_shaped_array_val(x) -> bool:
    import jax.core as jc
    return isinstance(jc.get_aval(x), jc.ShapedArray)

try:
    y = with_memory_space_constraint(x, memory_space=ms)
except NotImplementedError as e:
    if 'only supports arrays' in str(e):
        y = x  # skip constraint for non-array values
    else:
        raise

When it happens

Trigger: Passing a token or non-ShapedArray tracer (e.g. an effect token, or output of another exotic primitive) through with_memory_space_constraint during tracing; using it on values produced by ops whose avals are not ShapedArray.

Common situations: Composing Pallas kernels with stateful APIs that thread tokens; jax version changes introducing new aval types; annotated user-defined tracers flowing into memory-space constraint logic.

Related errors

• Mismatched type: {a, t}
• Cannot interpret value of type {typ} as an abstract array; i
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• input type mismatch for {_prim}
• at {keystr(path)}, got fwd output type {ty.str_short()} whic

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