{"record":{"id":"959cf8b488f75954","repo":"jax-ml/jax","slug":"eager-shard-map-cannot-return-a-jax-ref-please","errorCode":null,"errorMessage":"Eager shard_map cannot return a `jax.Ref`. Please wrap your shard_map in `jax.jit`.","messagePattern":"Eager shard_map cannot return a `jax\\.Ref`\\. Please wrap your shard_map in `jax\\.jit`\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1536,"sourceCode":"    return tree_map(_add_singleton, outs)\n  out_specs = list(out_specs) if type(out_specs) is tuple else out_specs\n  return shard_map(apply, mesh=concrete_mesh, in_specs=in_specs,\n                   out_specs=out_specs, check_vma=check_vma,\n                   axis_names=manual_axes)(*args)\n\neager_rules: dict[core.Primitive, Callable] = {}\n\ndef _device_put_eager_rule(mesh, *xs, srcs, devices, copy_semantics):\n  del mesh, srcs, copy_semantics\n  for device in devices:\n    if device is not None:\n      raise ValueError(\"device_put with explicit device not allowed within \"\n                       f\"shard_map-decorated functions, but got device {device}\")\n  return xs\neager_rules[dispatch.device_put_p] = _device_put_eager_rule\n\ndef _ref_raise_valueerror(*args, **kwargs):\n  raise ValueError(\n      \"Eager shard_map cannot return a `jax.Ref`. Please wrap\"\n      \" your shard_map in `jax.jit`.\")\n\neager_rules[core.ref_p] = _ref_raise_valueerror\neager_rules[core.empty_ref_p] = _ref_raise_valueerror\n\n# Batching\n\ndef used_axis_names(spec):\n  return _spec_to_mat(spec).vur\n\ndef _shard_map_batch(\n    trace: batching.BatchTrace, prim: core.Primitive, fun: Callable,\n    in_tracers: Sequence[batching.BatchTracer], mesh: Mesh,\n    in_specs, check_vma: bool, newly_manual_axes: frozenset,\n    debug_info) -> Sequence[batching.BatchTracer]:\n  in_vals, in_dims = unzip2(map(trace.to_batch_info, in_tracers))\n  spmd_axis_name = trace.axis_data.spmd_name","sourceCodeStart":1518,"sourceCodeEnd":1554,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1518-L1554","documentation":"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.","triggerScenarios":"Calling a shard_map function eagerly (without jax.jit) whose body creates or returns Refs (jax.experimental.ref / state-style code).","commonSituations":"Prototyping in notebooks with the Ref-based API; calling an exported step function without wrapping it in jit.","solutions":["Wrap the shard_map call in jax.jit","Refactor the body to return plain arrays instead of Refs if eager execution is needed","Call the function through a jit entry point (e.g. a jitted train_step)"],"exampleFix":"# before\nstep = shard_map(update_ref_body, mesh, in_specs=P('i'))\nstep(state, x)  # eager -> ValueError\n# after\nstep = jax.jit(shard_map(update_ref_body, mesh, in_specs=P('i')))\nstep(state, x)","handlingStrategy":"validation","validationCode":"def is_jitted() -> bool:\n    return not jax.config.jax_eager  # heuristic; better: ensure call site is under jax.jit\n# robust: wrap entry points with jax.jit at module definition time","typeGuard":null,"tryCatchPattern":"try: f(state, x) except ValueError as e: if 'jax.Ref' in str(e): f = jax.jit(f); f(state, x); else: raise","preventionTips":["Make jitted wrappers the only public entry points for Ref-based shard_map functions","Keep notebook experimentation behind a jitted step function"],"tags":["shard-map","ref","jit-required","jax"],"backgroundTag":"jit-required-op-called-eagerly","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}