Unbatch a tree

[limbryan]

Hi,

I am trying to unbatch a tree into a list of trees with the same structure of the length of the batch (more particularly for a flax NN structure). My attempt at this works but is really slow and time consuming and does not scale with the batch size. You can find the function I currently use below:

def single_tree_to_list_of_trees(tree, batch_size):
    '''
    single tree is in a batch form
    returns a list of trees with the params
    '''
    
    new_tree = [jax.tree_map(lambda x: x[i,:], tree) for i in range(batch_size)]

    return new_tree

Would appreciate if you could point me in the right direction or an implementation that would do this much more efficiently. I am probably just missing some functionality of jax that I am not aware of that allows me to do this easily. Thanks!

[cgarciae]

Hey @limbryan,

1. Is it slow to jit or slow after jit? These Python operations usually only run once per-jit so its not that bad.
2. What do you want to do with this list of trees? Can you frame your problem in terms of vmap?

[limbryan]

Hey @cgarciae,

Thanks for the quick reply and some suggestions.

1. I actually did not jit these functions before. But I have just tried after jit as well. It actually seems slower after. So in summary, it seems to be slow both ways.
2. I would like to have it in this data structure as a I would like to save some of the policies depending on certain characteristics of the corresponding policies in the batch. In summary, I would like to save them as individual policies and only some part of the batch and not all.

Hope this helps clear things up