Speed up mapping over single pytree

[c-weis]

Hello JAX people!

I'm trying to write a simulation using JAX. For outside design reasons, I'm tied to a hierarchical data structure like this for my model state:

ModuleTop
|- componentA
|- componentB
|- componentC
|- ModuleX
   |-componentD
   |-componentE
|- ModuleY
   |- componentF
   |- ModuleZ
      |-componentG

I would ideally like this state to be in the shape of a pytree and use JAX's features to perform operations on it. I coded up an example using tree.map below to show what I'm trying to achieve. The actual use case here is that I want to use Diffrax to solve a differential equation encoding a physics model - the map will be used to turn the state pytree into the derivative pytree.

My main question is:
Can I use vmap here to make speed gains?

from __future__ import annotations
from dataclasses import dataclass

import jax
import equinox as eqx


@dataclass
class Component:
    metadata: str
    data: float


class Module(eqx.Module):
    components: list[Component]
    submodules: list[Module]


def get_single_level_module():
    return Module(
        components=[
            Component(
                metadata="{name: 'A'}",
                data=0.5
            ),
            Component(
                metadata="{name: 'B'}",
                data=-1.0
            ),
        ],
        submodules=[]
    )


def test_map_nested_module():
    nested_module = get_single_level_module()
    nested_module.submodules.append(get_single_level_module())

    print("Pre-processed module:")
    print(nested_module)

    def double_data(component: Component) -> Component:
        return Component(
            metadata=component.metadata,
            data=component.data*2
        )

    processed_module = jax.tree.map(
        double_data,
        nested_module,
        is_leaf=lambda node: isinstance(node, Component)
    )

    print("Processed module:")
    print(processed_module)


if __name__ == "__main__":
    test_map_nested_module()

[jakevdp]

vmap operates on arrays of values within pytrees, and I don't see any arrays of values in your pytrees here, so vmap is probably not applicable.

[c-weis]

Oh thank you!

I guess what I'm really trying to achieve is to let JAX do the heavy lifting of flattening and un-flattening under the hood and quickly. If you have any pointers on what other option you recommend, I would really appreciate that. But thank you already for setting me free from my vmap-quest :)