Custom pytree strange leakage behavior

[eelregit]

I found that a custom pytree (of standard mutable pytree like list/dict)
behave strangely in the following example, where the treedef changes inside a function:

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class C:
    def __init__(self, l=[]):
        self.l = l

    def tree_flatten(self):
        return self.l, None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(children)
    
def f(c):
    c.l.append(0)
    return c

print(vars(C()))
print(vars(C()), vars(f(C())))

returns

{'l': []}
{'l': [0]} {'l': [0]}

The two C() return different outputs and the second one seem to be polluted by the f(C(0)) call.

[eelregit]

Another example is DifferentiableUniverseInitiative/jax_cosmo#81 (comment)
where treedef doesn't change but aux_data leak among instances.

[hawkinsp]

Isn't the problem here that the default value you are using in your __init__ is mutable? That's a known Python antipattern.

Try something like:

    def __init__(self, l=None):
        self.l = [] if l is None else l

[eelregit]

Thanks @hawkinsp! That explains it.