Easy way to inflate a Pytree to prefix for tree_map?

[femtomc]

Hi all!

I'm attempting to solve a Pytree filtering problem -- where I want to use one tree as a mask for another tree.

I want to support functionality where a user can partially specify the first tree (including a tree which is non-prefix!) -- for example:

d1 = {"y" : True, "z" : True}
d2 = {"y" : DeviceArray(...), "z" : DeviceArray(...), "q" : DeviceArray(...)}

Naive application of tree_map here will throw prefix errors -- I want to override the behavior such that tree_map with d1 specified as the structure will filter out "q" above, resulting in:

res = jtu.tree_map(f, d1, d2, is_leaf = ...)
# res = {"y" : mapped(DeviceArray(...)), "z" : mapped(DeviceArray(...))}

To complicate this problem even more, I want this to work on nested structures...

d1 = {"q": {"q": True}}
d2 = {"y": ..., "z": ..., "q" : {"q": DeviceArray(...), "z": DeviceArray(...)}

Critically, I think this means that the problem must depend on structure, and not just things like the length of the flattened Pytree leave array.

One other comment:

I think it's okay if the resulting Pytree has the same structure as the base Pytree that I am filtering.

This simplifies the problem to just being able to take aPytree and attempt to inflate it to a target Pytree shape. If the filter Pytree is not a prefix of the target Pytree, the resulting Pytree after filtering will be the same shape as the target Pytree, with all None. Otherwise, it will have entries for all prefixes that come from the filter Pytree.