jax-ml/jax · error · ValueError

devices already specified

Error message

devices already specified

What it means

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.

Source

Thrown at jax/experimental/colocated_python/func.py:102

    if out_specs_fn is None:
      out_specs_fn = self.out_specs_fn
    elif self.out_specs_fn is not None:
      raise ValueError("out_specs_fn already specified")

    if out_specs_treedef is None:
      out_specs_treedef = self.out_specs_treedef
    elif self.out_specs_treedef is not None:
      raise ValueError("out_specs already specified")
    if out_specs_leaves is None:
      out_specs_leaves = self.out_specs_leaves
    elif self.out_specs_leaves is not None:
      raise ValueError("out_specs already specified")

    if devices is None:
      devices = self.devices
    elif self.devices is not None:
      raise ValueError("devices already specified")
    elif not isinstance(devices, xc.DeviceList):
      devices = xc.DeviceList(tuple(devices))

    return Specialization(
        in_specs_treedef,
        in_specs_leaves,
        out_specs_fn,
        out_specs_treedef,
        out_specs_leaves,
        devices,
    )


def _get_spec(x: Any) -> api.ShapeDtypeStruct:
  """Extracts a spec for a value, which must be a JAX Array."""
  # TODO(hyeontaek): Allow Python values and automatically apply `shard_arg`
  # with a suitable sharding and layout.
  if not isinstance(x, jax.Array):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the second devices call; specialize separate function objects per device list
2. Derive the base (un-specialized) function and call .devices(...) once on each branch
3. If devices should follow the inputs, omit devices entirely and let it be inferred from argument sharding

Example fix

# before
f_cpu = f.devices(cpu_devs).devices(gpu_devs)  # ValueError
# after
f_cpu = f.devices(cpu_devs)
f_gpu = f.devices(gpu_devs)

def with_devices(f, devs):
    if getattr(f, 'devices', None) is not None:
        return f
    return f.devices(devs)

def has_devices(specialization) -> bool:
    return getattr(specialization, 'devices', None) is not None

try:
    f2 = f.devices(new_devs)
except ValueError as e:
    if 'devices already specified' in str(e):
        f2 = base.devices(new_devs)
    else:
        raise

When it happens

Trigger: Calling f.devices(devs_a).devices(devs_b), or update(devices=...) after devices were already given via specialize(devices=...).

Common situations: 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.

Related errors

• out_specs_fn already specified
• out_specs already specified
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• check_error takes an Error as argument, got type {type(error

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