Improve treedef

[janosg]

Problem

The JAX treedef contains more information. Therefore, their tree_unflatten function does not need an is_leaf argument. The same way we could make the registry argument unnecessary.

This is not just more convenient. It reduces error potential because the unflattening will automatically use the same options as the flattening.

To-Do

• Come up with better Treedef (namedtuple or Class) and implement it
• Adjust interface of tree_unflatten.

[janosg]

The usage of treedef in tree_transpose shows a few useful attributes that we should also have.