Create Custom PyTreeDef objects, or adapt existing ones.

[peterroelants]

Is it possible to create custom PyTreeDef objects, or adapting the structure of existing ones?

Context:
I'm stacking custom PyTree objects using jax.tree_map(lambda *xs: jnp.stack(xs, axis=0), *pytrees), however, if these objects contain metadata that are not leaf nodes, but are part of the PyTreeDef this fails.

The following example based on Treeo fails:

from dataclasses import dataclass

import jax
import jax.numpy as jnp
import treeo as to

@dataclass
class Person(to.Tree):
    height: jnp.array = to.field(node=True) # I am a node field!
    age_static: jnp.array = to.field(node=False) # I am a static field!, I should not be updated.
    name: str = to.field(node=False) # I am a static field!

persons = [
    Person(height=jnp.array(1.8), age_static=jnp.array(25.), name="John"),
    Person(height=jnp.array(1.7), age_static=jnp.array(100.), name="Wald"),
    Person(height=jnp.array(2.1), age_static=jnp.array(50.), name="Karen")
]

# Stack (struct of arrays instead of list of structs)
jax.tree_map(lambda *values: jnp.stack(values, axis=0), *persons)

This code results in this exception (open dropdown to see).

See also: cgarciae/treeo#23

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 18
     11     name: str = to.field(node=False) # I am a static field!
     13 persons = [
     14     Person(height=jnp.array(1.8), age_static=jnp.array(25.), name="John"),
     15     Person(height=jnp.array(1.7), age_static=jnp.array(100.), name="Wald"),
     16     Person(height=jnp.array(2.1), age_static=jnp.array(50.), name="Karen")
     17 ]
---> 18 jax.tree_map(lambda *values: jnp.stack(values, axis=0), *persons)

File ~/workspace/lcms_polymer_model/env/env_conda_local/lcms_polymer_model_env/lib/python3.10/site-packages/jax/_src/tree_util.py:199, in tree_map(f, tree, is_leaf, *rest)
    166 """Maps a multi-input function over pytree args to produce a new pytree.
    167 
    168 Args:
   (...)
    196   [[5, 7, 9], [6, 1, 2]]
    197 """
    198 leaves, treedef = tree_flatten(tree, is_leaf)
--> 199 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    200 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

File ~/workspace/lcms_polymer_model/env/env_conda_local/lcms_polymer_model_env/lib/python3.10/site-packages/jax/_src/tree_util.py:199, in <listcomp>(.0)
    166 """Maps a multi-input function over pytree args to produce a new pytree.
    167 
    168 Args:
   (...)
    196   [[5, 7, 9], [6, 1, 2]]
    197 """
    198 leaves, treedef = tree_flatten(tree, is_leaf)
--> 199 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    200 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

ValueError: Mismatch custom node data: {'_field_metadata': {'height': <treeo.types.FieldMetadata object at 0x7fb8b898ba00>, 'age_static': <treeo.types.FieldMetadata object at 0x7fb8b90c0a90>, 'name': <treeo.types.FieldMetadata object at 0x7fb8b8bf9db0>, '_field_metadata': <treeo.types.FieldMetadata object at 0x7fb8b89b56f0>, '_factory_fields': <treeo.types.FieldMetadata object at 0x7fb8b89b5750>, '_default_field_values': <treeo.types.FieldMetadata object at 0x7fb8b89b5660>, '_subtrees': <treeo.types.FieldMetadata object at 0x7fb8b89b5720>}, 'age_static': DeviceArray(25., dtype=float32, weak_type=True), 'name': 'John'} != {'_field_metadata': {'height': <treeo.types.FieldMetadata object at 0x7fb8b898ba00>, 'age_static': <treeo.types.FieldMetadata object at 0x7fb8b90c0a90>, 'name': <treeo.types.FieldMetadata object at 0x7fb8b8bf9db0>, '_field_metadata': <treeo.types.FieldMetadata object at 0x7fb8b89b56f0>, '_factory_fields': <treeo.types.FieldMetadata object at 0x7fb8b89b5750>, '_default_field_values': <treeo.types.FieldMetadata object at 0x7fb8b89b5660>, '_subtrees': <treeo.types.FieldMetadata object at 0x7fb8b89b5720>}, 'age_static': DeviceArray(100., dtype=float32, weak_type=True), 'name': 'Wald'}; value: Person(height=DeviceArray(1.7, dtype=float32, weak_type=True), age_static=DeviceArray(100., dtype=float32, weak_type=True), name='Wald').

I think it's clear to me how I'm getting this exception, it's because the stack method is not applied to the static fields (node=False). These static fields are stored in the PyTreeDef when flattening the tree.

I could do something like the following:

def stack_pytrees(pytrees):
    tree, *rest = pytrees
    leaves, treedefs = zip(*[jax.tree_util.tree_flatten(tree) for tree in pytrees])
    return treedefs[0].unflatten(jnp.stack(xs, axis=0) for xs in zip(*leaves))


stack_pytrees(persons)
# Person(height=DeviceArray([1.8, 1.7, 2.1], dtype=float32, weak_type=True), age_static=DeviceArray(25., dtype=float32, weak_type=True), name='John')

But this removes all metadata from all but the first PyTree. Ideally I would write some custom code mapping the PyTreeDef metadata to a list of that metadata or something. However, I cannot seem to figure out how to create a new custom PyTreeDef object where this is the case, or adapt the previous object. Can anyone help me trying to solve this problem (or tell me why I should not be doing this, and should solve this some other way)?

[jakevdp]

I've not heard of the treeo package, but for this to work on the JAX side I believe you'll need to register your class as a pytree. There is information on this at https://jax.readthedocs.io/en/latest/pytrees.html#extending-pytrees

There has been discussion surrounding automatic registration of dataclasses as pytrees (see #2371) but currently that registration is not automatic, so you need to use these registration mechanisms.

ALternatively, if you believe this is something that is supposed to happen automatically when you import treeo, I'd suggest inquiring in the discussions for that package.

[peterroelants]

The pytree ilnfrastructure has no mechansim for combining static metadata when combining trees.

So there is no way to manually create or edit objects of type PyTreeDef?

[jakevdp]

I don't know what you mean by "manually create or edit objects of type PyTreeDef".

[peterroelants]

I don't know what you mean by "manually create or edit objects of type PyTreeDef".

So when flattening I get a tuple of (leaves, treedefs) of which the leaves are JAX objects and treedefs are PyTreeDef objects:

pytrees = [
    RegisteredSpecial(x=1., y=2.),
    RegisteredSpecial(x=3., y=6.),
    RegisteredSpecial(x=4., y=8.)
]

leaves, treedefs = zip(*[jax.tree_util.tree_flatten(tree) for tree in pytrees])
print(f"{leaves=}")  # leaves=([1.0], [3.0], [4.0])
print(f"{treedefs=}")  # treedefs=(PyTreeDef(CustomNode(RegisteredSpecial[2.0], [*])), PyTreeDef(CustomNode(RegisteredSpecial[6.0], [*])), PyTreeDef(CustomNode(RegisteredSpecial[8.0], [*])))
print(f"{treedefs[0]}")  # PyTreeDef(CustomNode(RegisteredSpecial[2.0], [*]))
print(f"{type(treedefs[0])}")  # <class 'jaxlib.xla_extension.pytree.PyTreeDef'>

It's these jaxlib.xla_extension.pytree.PyTreeDef that I want to work with and edit or create.

[jakevdp]

Oh, OK. As far as I know, the only supported mechanism for creating PyTree objects is by using tree_util.tree_flatten. Perhaps you can use that in the "manual" creation process that you have in mind?

[peterroelants]

Oh, OK. As far as I know, the only supported mechanism for creating PyTree objects is by using tree_util.tree_flatten.

I see. I also found the definition of PyTreeDef in the TensorFlow repo and I noticed that indeed no constructor is defined, nor any other methods that would be helpful to me.

Thanks for your help, and sorry for any confusion initially, I could have been clearer in hindsight.

[peterroelants]

From the discussion at: #13768 (reply in thread)

It is not possible to create or edit PyTreeDef objects.

jaxlib.xla_extension.pytree.PyTreeDef are defined as C++ code in the Tensorflow XLA Library. These do not seem to expose methods that could be useful in solving this issue.