jax-ml/jax · error · ValueError

out_specs already specified

Error message

out_specs already specified

What it means

Raised when out_specs_treedef is supplied to update() but the specialization already carries an out_specs_treedef (i.e. out_specs was previously set). Internally out_specs is flattened into (treedef, leaves), and neither part may be reassigned after being set. It guards against conflicting output specifications in the builder chain.

Source

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

    """Creates a new specialization with overrides."""
    if in_specs_treedef is None:
      in_specs_treedef = self.in_specs_treedef
    elif self.in_specs_treedef is not None:
      raise ValueError("in_specs already specified")
    if in_specs_leaves is None:
      in_specs_leaves = self.in_specs_leaves
    elif self.in_specs_leaves is not None:
      raise ValueError("in_specs already specified")

    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Delete the redundant out_specs argument from the later call
2. Create a new base specialization if you need different output specs
3. Check whether a decorator/wrapper already applied out_specs before your call

Example fix

# before
spec = f.update(out_specs=s1).update(out_specs=s2)  # ValueError
# after
spec_a = f.update(out_specs=s1)
spec_b = f.update(out_specs=s2)

def set_out_specs_once(f, out_specs):
    if f.in_specs_leaves is not None and f.out_specs_treedef is not None:
        return f  # already fully specified
    return f.update(out_specs=out_specs)

def out_specs_set(specialization) -> bool:
    return getattr(specialization, 'out_specs_treedef', None) is not None

try:
    f2 = f.update(out_specs=s)
except ValueError as e:
    if 'out_specs already specified' in str(e):
        pass  # keep existing spec
    else:
        raise

When it happens

Trigger: Calling update()/specialize() with out_specs=... (or its flattened treedef) after out_specs was already specified earlier in the chain, e.g. f.out_specs(s1).out_specs(s2).

Common situations: Reusing a partially configured colocated function object and adding output specs again; copy-paste of configuration blocks that both set out_specs.

Related errors

• out_specs_fn already specified
• devices 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/3c25dfda7dccb735. Report an issue: GitHub.