{"record":{"id":"39b07e4b240f1bc2","repo":"jax-ml/jax","slug":"device-put-with-explicit-device-not-allowed-within","errorCode":null,"errorMessage":"device_put with explicit device not allowed within shard_map-decorated functions, but got device {device}","messagePattern":"device_put with explicit device not allowed within shard_map-decorated functions, but got device (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1530,"sourceCode":"  __repr__ = __str__  # for debuggers, like `p x`\n\ndef _prim_applier(prim, check_vma, params_tup, concrete_mesh, manual_axes,\n                  in_specs, out_specs, *args):\n  def apply(*args):\n    outs = prim.bind(*map(_rem_singleton, args), **dict(params_tup))\n    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(","sourceCodeStart":1512,"sourceCodeEnd":1548,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1512-L1548","documentation":"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.","triggerScenarios":"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.","commonSituations":"Copy-pasting host-to-device placement helper code into a sharded body; libraries that eagerly place buffers on device 0.","solutions":["Remove the explicit device argument; use jax.device_put(x) without a device or plain array ops","Move the device_put outside the shard_map-wrapped function","Rely on the shard_map in_specs/out_specs and mesh to control placement"],"exampleFix":"# before\ng = shard_map(lambda a: jax.device_put(a, jax.devices()[0]), mesh, ...)\n# after\ng = shard_map(lambda a: a * 2, mesh, ...)  # placement handled by shard_map","handlingStrategy":"validation","validationCode":"import jax._src.dispatch as dispatch\ndef no_explicit_devices(fn_src): ...  # static check\n# simplest guard: audit body for jax.device_put calls with a device argument","typeGuard":null,"tryCatchPattern":"try: f(x) except ValueError as e: if 'device_put with explicit device' in str(e): strip device args from body; else: raise","preventionTips":["Never place arrays on explicit devices inside shard_map bodies","Let in_specs/out_specs control placement"],"tags":["shard-map","device-put","placement","jax"],"backgroundTag":"explicit-device-in-sharded-region","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}