jax-ml/jax · error · ValueError

out_specs_fn already specified

Error message

out_specs_fn already specified

What it means

Raised by ColocatedFunction.update when out_specs_fn is passed to a specialization that already has an out_specs_fn set. ColocatedPython accumulates in_specs/out_specs/devices across update/specialize calls, and each output-spec slot can only be set once. The error prevents ambiguous conflicting output specifications.

Source

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

      out_specs_fn: Callable[..., ShapeDtypeStructTree] | None = None,
      out_specs_treedef: tree_util.PyTreeDef | None = None,
      out_specs_leaves: tuple[api.ShapeDtypeStruct, ...] | None = None,
      devices: Sequence[jax.Device] | xc.DeviceList | None = None,
  ):
    """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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the second out_specs_fn assignment; each spec may be set only once per specialization chain
2. If overriding is intended, build a fresh specialization from the original function instead of updating the existing one
3. Audit helper functions/wrappers that set out_specs_fn on your behalf before you set it again

Example fix

# before
f2 = f.update(out_specs_fn=fn1).update(out_specs_fn=fn2)  # ValueError
# after
f2 = f.update(out_specs_fn=fn1)
f3 = f.update(out_specs_fn=fn2)

def safe_out_specs_fn(f, fn):
    state = f.__self__ if hasattr(f, '__self__') else f
    # inspect current specialization before updating
    spec = getattr(state, '_specialization', None) or state
    if getattr(spec, 'out_specs_fn', None) is not None:
        raise RuntimeError('out_specs_fn already set; refusing to overwrite')
    return f.update(out_specs_fn=fn)

def has_out_specs_fn(specialization) -> bool:
    return getattr(specialization, 'out_specs_fn', None) is not None

try:
    f2 = f.update(out_specs_fn=fn)
except ValueError as e:
    if 'out_specs_fn already specified' in str(e):
        f2 = f  # already configured; reuse existing
    else:
        raise

When it happens

Trigger: Calling f.out_specs(fn1).out_specs(fn2), or specialize()/update() with out_specs_fn=... after out_specs_fn was already provided (directly or via out_specs=...).

Common situations: Chaining builder-style configuration on a colocated_python function, e.g. setting output specs in a helper and then again at the call site; refactoring that splits spec configuration across multiple functions.

Related errors

• out_specs 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/bebd6478faf76db0. Report an issue: GitHub.