making use of aux data of custom Pytree node in tree map

[SaitejaUtpala]

Consider following custom Pytree with auxillary data.

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class CustomPytreeWithAuxData:
    def __init__(self, value, attr) :
        self.value = value
        self.attr = attr

   
    def tree_flatten(self):
        children = (self.value,)
        aux_data = self.attr
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(children[0], aux_data)

Now consider transforming a two Pytree which have node of above type using jax.tree_map. But I want to transform using attribute which is an aux data, which will not be available if we using treemap, as transformation is done using only leaves.

jax.tree_map( lambda leave1, leave2 : # auxdata is not accessible here, Pytree1, Pytree2)

Any best solution ?

Thanks in advance.

[SaitejaUtpala]

cleared in #12711