TF Nest alternative

[albertz]

For some code, it is useful to have functions like what TF nest (from tensorflow.python.util import nest) provides, such as nest.map_structure, nest.flatten, nest.pack_sequence_as, and others.

So far, it was not really a problem to use just the TF nest, as we anyway have TF as dependency (although for RC, this was debatable: rwth-i6/returnn_common#27).

Now, as we support different backends again (#1120), we want to have a framework-independent solution.

So, what are the alternatives:

• Create a simple implementation ourselves. It's not really so hard. But it's kind of reinventing the wheel.
• Copy over some existing implementation (e.g. TF nest) and keep our own copy.
• Use some other alternative package, i.e. add a new dependency.

(Similar discussion here: rwth-i6/returnn_common#27 (comment))

Some known implementations:

• TF nest. No own package, thus depends on whole TF.
• Pytrees from JAX. No own package, depends on whole JAX.
• Tree from DeepMind. Independent package. Last commit Jan 2023. pip install dm-tree
• PyTorch PyTree (torch.utils._pytree). No own package, thus depends on whole PyTorch. Also slow, hidden, but might be made public, and speed might improve (torch.utils._pytree -> stable pytorch/pytorch#65761).
• OpTree from metaopt. Looks good. pip install optree.
• iteration-utilities (GitHub 60 stars). Last commit Dec 2020. So basically unmaintained.
• more-itertools (GitHub >2k stars). Last commit Apr 2023.

I did not really looked into them in detail whether they really exactly provide what we need (which is mostly exactly what TF nest provides).

[albertz]

I'm currently leaning towards those two options:

• Copy over TF nest to RETURNN. Just because I'm familiar with it.
• Use Tree from DeepMind. Seems well-maintained and exactly what we need (in contrast to the other packages which provide other stuff). But this would add a new dependency.

[curufinwe]

From these last 2 options I think I would prefer to add the dependency to tree. This will keep the RETURNN code small and more approachable to new developers.

[albertz]

Ok then it's decided: we use Tree from DeepMind.

[albertz]

I also adapted https://github.com/rwth-i6/returnn/wiki/PyCharm-Configuration.

[albertz]

Note that this is currently only used (and needed) for the RF (returnn.frontend).

Edit Ok, the RF is also imported when the engine is selected in RETURNN, to initialize and set a reasonable RF backend. This is done always at startup. I'm not sure if we want to avoid this in any way. It would probably get ugly. But so it means everyone really needs to install dm-tree now. But maybe it's not really such a big problem?

[albertz]

Getting back to this, for completeness, some new alternatives:

• PyTorch PyTree: torch.utils._pytree, part of PyTorch, slow, hidden, but might be made public, and speed might improve (torch.utils._pytree -> stable pytorch/pytorch#65761)
• OpTree: https://github.com/metaopt/optree