What causes "out_specs already specified" in jax-ml/jax?

Delete the redundant out_specs argument from the later call

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.

Solutions

Delete the redundant out_specs argument from the later call
Create a new base specialization if you need different output specs
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)

Defensive patterns

Strategy: validation

Validate before calling

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)

Type guard

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

Try / catch

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

Prevention

Centralize out_specs assignment in one place
Branch to a new base specialization for alternate output specs
Log which component sets each spec slot once

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

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

Appendix: 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)