Convert pytree of arrays to list of pytress without leading dim

[alonfnt]

Hi, I was looking for a way to convert a namedtupple where all the components have the same leading dimension (after being used on a tree_map) to a list of that same namedtuple, so that I can iterate over them.

Something like this:

from collections import namedtuple
import jax
import jax.numpy as jnp


Foo = namedtuple('Foo', 'a b')

foo = Foo(a=jnp.ones((10, 5, 2)), b=jnp.ones((10, 1)))

# Something equivalent to this where to output should be a list of length 10:
foo_list = [Foo(a=a, b=b) for a, b in zip(foo.a, foo.b)]

I was looking into jax.tree_util.tree_transform but it does not seem to be what I need (I think).

[jakevdp]

I think list(map(Foo, *foo)) should do what you want for the case where Foo is a simple namedtuple containing arrays.

[alonfnt]

That is nice. I'm not sure why I was thinking on pytrees instead of standard python 😓. Thanks