{"record":{"id":"274082893f7687db","repo":"jax-ml/jax","slug":"devices-already-specified","errorCode":null,"errorMessage":"devices already specified","messagePattern":"devices already specified","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/colocated_python/func.py","lineNumber":102,"sourceCode":"\n    if out_specs_fn is None:\n      out_specs_fn = self.out_specs_fn\n    elif self.out_specs_fn is not None:\n      raise ValueError(\"out_specs_fn already specified\")\n\n    if out_specs_treedef is None:\n      out_specs_treedef = self.out_specs_treedef\n    elif self.out_specs_treedef is not None:\n      raise ValueError(\"out_specs already specified\")\n    if out_specs_leaves is None:\n      out_specs_leaves = self.out_specs_leaves\n    elif self.out_specs_leaves is not None:\n      raise ValueError(\"out_specs already specified\")\n\n    if devices is None:\n      devices = self.devices\n    elif self.devices is not None:\n      raise ValueError(\"devices already specified\")\n    elif not isinstance(devices, xc.DeviceList):\n      devices = xc.DeviceList(tuple(devices))\n\n    return Specialization(\n        in_specs_treedef,\n        in_specs_leaves,\n        out_specs_fn,\n        out_specs_treedef,\n        out_specs_leaves,\n        devices,\n    )\n\n\ndef _get_spec(x: Any) -> api.ShapeDtypeStruct:\n  \"\"\"Extracts a spec for a value, which must be a JAX Array.\"\"\"\n  # TODO(hyeontaek): Allow Python values and automatically apply `shard_arg`\n  # with a suitable sharding and layout.\n  if not isinstance(x, jax.Array):","sourceCodeStart":84,"sourceCodeEnd":120,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/colocated_python/func.py#L84-L120","documentation":"Raised by ColocatedFunction.update when a devices argument is passed but the specialization already has devices set. ColocatedPython pins a function to one device list per specialization; re-specifying it would be ambiguous. Devices can alternatively be inferred from argument sharding at call time.","triggerScenarios":"Calling f.devices(devs_a).devices(devs_b), or update(devices=...) after devices were already given via specialize(devices=...).","commonSituations":"Setting devices in a setup helper and again per-call; switching device lists by mutating a shared specialized function instead of creating a new specialization.","solutions":["Remove the second devices call; specialize separate function objects per device list","Derive the base (un-specialized) function and call .devices(...) once on each branch","If devices should follow the inputs, omit devices entirely and let it be inferred from argument sharding"],"exampleFix":"# before\nf_cpu = f.devices(cpu_devs).devices(gpu_devs)  # ValueError\n# after\nf_cpu = f.devices(cpu_devs)\nf_gpu = f.devices(gpu_devs)","handlingStrategy":"validation","validationCode":"def with_devices(f, devs):\n    if getattr(f, 'devices', None) is not None:\n        return f\n    return f.devices(devs)","typeGuard":"def has_devices(specialization) -> bool:\n    return getattr(specialization, 'devices', None) is not None","tryCatchPattern":"try:\n    f2 = f.devices(new_devs)\nexcept ValueError as e:\n    if 'devices already specified' in str(e):\n        f2 = base.devices(new_devs)\n    else:\n        raise","preventionTips":["Specialize one function object per device list","Skip devices= entirely when you want sharding-based inference","Keep the un-configured base function around for re-specialization"],"tags":["jax","colocated-python","devices","api-misuse"],"backgroundTag":"duplicate-option-specification","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}