Add mapLeaves operation to Tree

[benikm91]

Add mapLeaves operation to Tree (similar to jax.tree_util.tree_leaves).

General Motivation: Add reduce operation over a TensorTree structure.
Specific Motivation: Calculate Gradient Norm for Gradient Clipping (GPT-2 in deepwit).

This PR adds mapLeaves to TensorTree and FloatTree, which is similar to jax.tree_util.tree_leaves from JAX.

Difference to JAX: jax.tree_util.tree_leaves just returns a List of Any (without mapping), which makes less sense in a strongly typed language. Therefore, we require a mapping function f: [T <: Tuple, V] => (Labels[T]) ?=> (Tensor[T, V] => A) to unify all leaves (tensors) to a shared type A. Reduction across A can then be done on the Iterator[A] object.

Note:

• We can't return a Tree of A's as the Tree Structure can be a case class of (different-shaped) tensors. So, e.g., Params[A] does not make sense in such cases. Therefore, we (must) return a flat Iterator over leaves.
• Returning an Iterator without a map function f is not possible, as Tensor[T, V] has to type parameters T and V, which differ across leaves. It would require an Iterator with Polymorphic functions in map, reduce, etc.

[benikm91]

Last commit implements #87 in new TensorTree and FloatTree.

[marcelluethi]

Thanks for the change. I suppose this supersedes #87, which then could be closed.