Can the children of a class registered as a pytree change dynamically?

[beckermr]

I have cooked up the following example code

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp

from jax.tree_util import register_pytree_node_class
import functools
import jax.numpy as jnp
from jax.tree_util import tree_flatten, tree_unflatten


def show_example(structured):
    flat, tree = tree_flatten(structured)
    unflattened = tree_unflatten(tree, flat)
    print(f"{structured=}\n  {flat=}\n  {tree=}\n  {unflattened=}")


def lazy_property(func):
    attname = func.__name__ + "_cached"

    @property
    @functools.wraps(func)
    def _func(self):
        if attname not in self._workspace:
            self._workspace[attname] = func(self)
        return self._workspace[attname]

    return _func


@register_pytree_node_class
class A:
    def __init__(self, _x):
        self._x = _x
        self._workspace = {}

    @property
    def x(self):
        return self._x

    @lazy_property
    def y(self):
        return self.x * self.x

    def tree_flatten(self):
        children = (self.x, self._workspace)
        aux_data = None
        return (children, aux_data)
    
    @classmethod
    def tree_unflatten(cls, aux_data, children):
        ret = cls(children[0])
        ret._workspace.update(children[1])
        return ret


@jax.jit
def get_y(a):
    return a.y

def fn(x):
    a = A(x)
    z = a.y**2
    return z + jnp.sin(get_y(a))


print("these calls all run properly:")
jfn = jax.jit(fn)
vfn = jax.vmap(fn)
jvfn = jax.jit(jax.vmap(fn))

gfn = jax.grad(fn)
vgfn = jax.vmap(jax.grad(fn))
jgfn = jax.jit(jax.grad(fn))
jvgfn = jax.jit(jax.vmap(jax.grad(fn)))

print(f"\t{fn(3.0)=}")
print(f"\t{jfn(3.0)=}")

print(f"\t{vfn(jnp.arange(10))[3]=}")
print(f"\t{jvfn(jnp.arange(10))[3]=}")

print(f"\t{gfn(3.0)=}")
print(f"\t{jgfn(3.0)=}")
eps = 1e-4
print("\tcheck grad=", (fn(3.0 + eps) - fn(3.0 - eps)) / 2.0 / eps)

print(f"\t{vgfn(jnp.arange(10, dtype=float))[3]=}")
print(f"\t{jvgfn(jnp.arange(10, dtype=float))[3]=}")


print("\n...but the children change shape:")
a = A(5.0)
show_example(a)
a.y
show_example(a)

It produces:

these calls all run properly:
	fn(3.0)=Array(81.41211849, dtype=float64, weak_type=True)
	jfn(3.0)=Array(81.41211849, dtype=float64, weak_type=True)
	vfn(jnp.arange(10))[3]=Array(81.41211849, dtype=float64)
	jvfn(jnp.arange(10))[3]=Array(81.41211849, dtype=float64)
	gfn(3.0)=Array(102.53321843, dtype=float64, weak_type=True)
	jgfn(3.0)=Array(102.53321843, dtype=float64, weak_type=True)
	check grad= 102.53321885215883
	vgfn(jnp.arange(10, dtype=float))[3]=Array(102.53321843, dtype=float64)
	jvgfn(jnp.arange(10, dtype=float))[3]=Array(102.53321843, dtype=float64)

...but the children change shape:
structured=<__main__.A object at 0x14eaea6b0>
  flat=[5.0]
  tree=PyTreeDef(CustomNode(A[None], [*, {}]))
  unflattened=<__main__.A object at 0x14eb6b850>
structured=<__main__.A object at 0x14eaea6b0>
  flat=[5.0, 25.0]
  tree=PyTreeDef(CustomNode(A[None], [*, {'y_cached': *}]))
  unflattened=<__main__.A object at 0x14ee6f940>

Everything appears to work but the children of the class change dynamically as it is used. Is this behavior allowed in JAX?

[jakevdp]

I think this is mostly OK – the one situation this might cause problems is if a function is if the pytree can be mutated during abstract evaluation of a function. This might, for example, mean that when you pre-compile a function using jax.jit(f).lower(dynamic_pytree).compile(), when you later call that function the input will no longer match the expected argument signature.

I can't think of other cases where it would be problematic, but they may exist – mutation can be tricky.

[beckermr]

Thank you @jakevdp! I've also seen a version of this (in jax-cosmo) where the workspace is discarded through the flattening process:

   def tree_flatten(self):
        children = (self.x,)
        aux_data = None
        return (children, aux_data)
    
    @classmethod
    def tree_unflatten(cls, aux_data, children):
        ret = cls(children[0])
        return ret

Is this version any better/worse (beyond increased cache misses)? Does it have other issues since now the class has a dynamic part that is not captured in the children?

[jakevdp]

I think that's probably fine – it's probably truer to what's actually happening. The issue with the cache-based attribute lookup is that if the cache is created within a transformed function, the cache won't persist outside it no matter how you structure the flattening/unflattening. Probably better to explicitly avoid flattening the cache.