jax-ml/jax · error · ValueError

The names should be exclusive and should not intersect in `n

Error message

The names should be exclusive and should not intersect in `names_which_can_be_saved` and `names_which_can_be_offloaded`. Got names_which_can_be_saved={set(names_which_can_be_saved)}, names_which_can_be_offloaded={set(names_which_can_be_offloaded)} and the intersection={set(intersection)}

What it means

While flattening with keypaths, jaxlib encounters a tuple whose type (or base) looks like a namedtuple and reads _fields to emit GetAttrKey entries. If _fields is not a tuple or its length differs from the tuple instance's size, this error is thrown during full tree flattening (not just one-level).

Source

Thrown at jax/_src/ad_checkpoint.py:174

  offload_dst: str

  def __call__(self, prim, *_, **params) -> Any:
    if prim is name_p and params['name'] in self.names_which_can_be_saved:
      return pe.Saveable
    if prim is name_p and params['name'] in self.names_which_can_be_offloaded:
      return pe.Offloadable(src=self.offload_src, dst=self.offload_dst)
    return pe.Recompute  # not saveable unless it's in the allow-list

def save_and_offload_only_these_names(
    *, names_which_can_be_saved, names_which_can_be_offloaded,
    offload_src, offload_dst):
  """Same as ``save_only_these_names``, but offload to CPU memory instead of
  recomputing."""
  names_which_can_be_saved = frozenset(names_which_can_be_saved)
  names_which_can_be_offloaded = frozenset(names_which_can_be_offloaded)
  intersection = names_which_can_be_saved & names_which_can_be_offloaded
  if intersection:
    raise ValueError(
        "The names should be exclusive and should not intersect in"
        " `names_which_can_be_saved` and `names_which_can_be_offloaded`. Got"
        f" names_which_can_be_saved={set(names_which_can_be_saved)},"
        f" names_which_can_be_offloaded={set(names_which_can_be_offloaded)} and"
        f" the intersection={set(intersection)}")
  return SaveAndOffloadOnlyTheseNames(
      names_which_can_be_saved, names_which_can_be_offloaded,
      offload_src, offload_dst)


def save_from_both_policies(policy_1, policy_2):
  """Logical OR of the given policies.

  A residual is saveable iff it is saveable according to either policy."""
  def policy(prim, *args, **params):
    out1 = policy_1(prim, *args, **params)
    out2 = policy_2(prim, *args, **params)
    if not (isinstance(out1, bool) and isinstance(out2, bool)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check type(node)._fields is a tuple and len(node) == len(type(node)._fields) for suspect nodes
2. Reconstruct objects as real namedtuples after deserialization
3. Avoid defining _fields on non-namedtuple tuple subclasses
4. Use is_leaf to short-circuit flattening of the malformed node

Example fix

# before
restored = pickle.loads(blob)  # namedtuple with fewer elements than _fields
jax.tree.flatten_with_path(restored)

# after
Restored = collections.namedtuple('Restored', Restored._fields[:len(restored)])
restored = Restored(*restored)
jax.tree.flatten_with_path(restored)

def validate_tree_for_keypaths(tree):
    for x in jax.tree.leaves(tree, is_leaf=lambda v: isinstance(v, tuple) and hasattr(type(v), '_fields')):
        pass  # leaf override keeps malformed namedtuples out of flattening

def namedtuple_intact(x) -> bool:
    f = getattr(type(x), '_fields', None)
    return not (isinstance(x, tuple) and f is not None) or (isinstance(f, tuple) and len(f) == len(x))

def is_safe_pytree_input(tree) -> bool:
    stack = [tree]
    while stack:
        v = stack.pop()
        if isinstance(v, tuple) and hasattr(type(v), '_fields'):
            if not isinstance(type(v)._fields, tuple) or len(type(v)._fields) != len(v):
                return False
        elif isinstance(v, (list, tuple)):
            stack.extend(v)
        elif isinstance(v, dict):
            stack.extend(v.values())
    return True

When it happens

Trigger: jax.tree.flatten_with_path / tree_map_with_path over a structure containing a pseudo-namedtuple: a tuple subclass with a bogus _fields attribute, or a namedtuple whose instance length was altered (e.g. via tuple.__new__ tricks or serialization that dropped fields).

Common situations: Deserializing namedtuple-like objects (pickle from an older class version with different fields), libraries faking namedtuples, duck-typed _fields properties that return lists or generators.

Related errors

• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• stop_gradient only works on valid JAX arrays, but input argu
• {self.__class__.__name__} has no attribute {name}
• Expected named tuple, got %s.

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