jax-ml/jax · error · ValueError

Eager shard_map cannot return a `jax.Ref`. Please wrap your

Error message

Eager shard_map cannot return a `jax.Ref`. Please wrap your shard_map in `jax.jit`.

What it means

Eager (non-jit) shard_map execution cannot produce Ref outputs (mutable buffers used by e.g. shard_map's in-place variant); Refs only make sense inside a jit boundary. The eager rule for ref_p/empty_ref_p raises this ValueError.

Source

Thrown at jax/_src/shard_map.py:1536

    return tree_map(_add_singleton, outs)
  out_specs = list(out_specs) if type(out_specs) is tuple else out_specs
  return shard_map(apply, mesh=concrete_mesh, in_specs=in_specs,
                   out_specs=out_specs, check_vma=check_vma,
                   axis_names=manual_axes)(*args)

eager_rules: dict[core.Primitive, Callable] = {}

def _device_put_eager_rule(mesh, *xs, srcs, devices, copy_semantics):
  del mesh, srcs, copy_semantics
  for device in devices:
    if device is not None:
      raise ValueError("device_put with explicit device not allowed within "
                       f"shard_map-decorated functions, but got device {device}")
  return xs
eager_rules[dispatch.device_put_p] = _device_put_eager_rule

def _ref_raise_valueerror(*args, **kwargs):
  raise ValueError(
      "Eager shard_map cannot return a `jax.Ref`. Please wrap"
      " your shard_map in `jax.jit`.")

eager_rules[core.ref_p] = _ref_raise_valueerror
eager_rules[core.empty_ref_p] = _ref_raise_valueerror

# Batching

def used_axis_names(spec):
  return _spec_to_mat(spec).vur

def _shard_map_batch(
    trace: batching.BatchTrace, prim: core.Primitive, fun: Callable,
    in_tracers: Sequence[batching.BatchTracer], mesh: Mesh,
    in_specs, check_vma: bool, newly_manual_axes: frozenset,
    debug_info) -> Sequence[batching.BatchTracer]:
  in_vals, in_dims = unzip2(map(trace.to_batch_info, in_tracers))
  spmd_axis_name = trace.axis_data.spmd_name

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the shard_map call in jax.jit
2. Refactor the body to return plain arrays instead of Refs if eager execution is needed
3. Call the function through a jit entry point (e.g. a jitted train_step)

Example fix

# before
step = shard_map(update_ref_body, mesh, in_specs=P('i'))
step(state, x)  # eager -> ValueError
# after
step = jax.jit(shard_map(update_ref_body, mesh, in_specs=P('i')))
step(state, x)

def is_jitted() -> bool:
    return not jax.config.jax_eager  # heuristic; better: ensure call site is under jax.jit
# robust: wrap entry points with jax.jit at module definition time

try: f(state, x) except ValueError as e: if 'jax.Ref' in str(e): f = jax.jit(f); f(state, x); else: raise

When it happens

Trigger: Calling a shard_map function eagerly (without jax.jit) whose body creates or returns Refs (jax.experimental.ref / state-style code).

Common situations: Prototyping in notebooks with the Ref-based API; calling an exported step function without wrapping it in jit.

Related errors

• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• Primitive {prim_name} requires varying manual axes to match,
• {} function carry input and carry output must have equal typ
• Scan carry input and output got mismatched varying manual ax

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