jax-ml/jax · error · AttributeError

{self.__class__.__name__} has no attribute {name}

Error message

{self.__class__.__name__} has no attribute {name}

What it means

FromPickle found a non-None node_data (t[2]) for a node kind that should carry no node data (leaf, tuple, list, etc.). The pickle's node encoding is inconsistent with the expected format.

Source

Thrown at jax/_src/ad_util.py:123

stop_gradient_p.def_abstract_eval(lambda x: x)


# User-facing version of `Zero`
class SymbolicZero:
  def __init__(self, aval: core.AbstractValue) -> None:
    self.aval = aval

  def __repr__(self) -> str:
    return self.__class__.__name__

  # TODO(mattjj,frostig): this forwards attr lookup to self.aval delegate;
  # should dedup with core.Tracer.__getattr__ which does the same thing
  def __getattr__(self, name):
    # if the aval property raises an AttributeError, gets caught here
    try:
      attr = getattr(self.aval, name)
    except KeyError as err:
      raise AttributeError(
          f"{self.__class__.__name__} has no attribute {name}"
      ) from err
    else:
      t = type(attr)
      if t is core.aval_property:
        return attr.fget(self)
      elif t is core.aval_method:
        return types.MethodType(attr.fun, self)
      else:
        return attr

  @staticmethod
  def from_primal_value(val: Any) -> SymbolicZero:
    return SymbolicZero(typeof(val).to_tangent_aval())

def zero_from_primal(val, symbolic_zeros=False):
  def f(x):
    t_aval = typeof(x).to_tangent_aval()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regenerate the treedef from actual Python objects instead of unpickling
2. Re-serialize with the matching jaxlib version
3. Validate pickle provenance (jaxlib version) before loading

try:
    treedef = pickle.loads(data)
except Exception as e:
    if 'Malformed pickled PyTreeDef' in str(e):
        _, treedef = jax.tree_util.tree_flatten(reference)
    else:
        raise

When it happens

Trigger: Unpickling a PyTreeDef where a non-custom node tuple has a non-None third element, typically from malformed or version-skewed pickles.

Common situations: Cross-version unpickling of saved treedefs; corrupted checkpoint files; third-party code that synthesizes pickled treedefs incorrectly.

Related errors

• The names should be exclusive and should not intersect in `n
• stop_gradient only works on valid JAX arrays, but input argu
• Could not find type: %s.
• Malformed pickled PyTreeDef, expected 2-tuple
• Malformed pickled SequenceKey, expected 1-tuple

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