jax-ml/jax · error · ValueError

device_put with explicit device not allowed within shard_map

Error message

device_put with explicit device not allowed within shard_map-decorated functions, but got device {device}

What it means

Within a shard_map-decorated function, device_put with an explicit target device is disallowed because placement decisions belong to the outer sharding, not the mapped body. Passing any non-None device raises this ValueError.

Source

Thrown at jax/_src/shard_map.py:1530

  __repr__ = __str__  # for debuggers, like `p x`

def _prim_applier(prim, check_vma, params_tup, concrete_mesh, manual_axes,
                  in_specs, out_specs, *args):
  def apply(*args):
    outs = prim.bind(*map(_rem_singleton, args), **dict(params_tup))
    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the explicit device argument; use jax.device_put(x) without a device or plain array ops
2. Move the device_put outside the shard_map-wrapped function
3. Rely on the shard_map in_specs/out_specs and mesh to control placement

Example fix

# before
g = shard_map(lambda a: jax.device_put(a, jax.devices()[0]), mesh, ...)
# after
g = shard_map(lambda a: a * 2, mesh, ...)  # placement handled by shard_map

import jax._src.dispatch as dispatch
def no_explicit_devices(fn_src): ...  # static check
# simplest guard: audit body for jax.device_put calls with a device argument

try: f(x) except ValueError as e: if 'device_put with explicit device' in str(e): strip device args from body; else: raise

When it happens

Trigger: Calling jax.device_put(x, device) or a library function that pins to a device (e.g. device_put(x, jax.devices()[0])) inside the shard_map body.

Common situations: Copy-pasting host-to-device placement helper code into a sharded body; libraries that eagerly place buffers on device 0.

Related errors

• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• The devices() method was called on {self._error_repr()}.{sel
• sharding with memory_kind is not allowed. Please use `jax.de
• Primitive {prim_name} requires varying manual axes to match,

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