Operations should allow shape broadcasting about tree-math HOT 5 OPEN

The intention of the code makes sense, however according to the current semantics of Vector the operation seems weird because v has shape (15,) but the after the norm per leaf it has shape (2,), and division between such vectors would make no sense.

Maybe we need a looser abstraction (e.g. Numeric) that just lets you do leaf wise math operations but doesn't have array-like semantics.

from tree-math.

I think I see. I was somehow expecting Vector to allow leaf-wise math. The abstraction you suggest could be nice to avoid tree_map calls in general. For the above example, having something like
normalized_tree = tree / tm.norm(tree), with tree an instance of Numeric may be ideal. One example use-case is implementing power iteration with a pytree input.

from tree-math.

Leaf-wise math sounds potentially useful but as @cgarciae notes it's definitely a different data model. I would love to see a more fully fleshed out use-case for this functionality.

from tree-math.

Since a lot of the operations from Vector are also leaf-wise there is the possibility that Numeric / NumericMixin might serve as a base class to Vector and friends. One of the big differences with Vector would be the dot method, it would either have to also be leaf-wise and return a pytree or just not included as an operation.

from tree-math.

My mental model for objects like this is that they are something closer to ragged arrays, e.g., a 2D matrix where the first dimension corresponds to the number of pytree leaves, and the second dimension corresponds to the (variable) size of each leaf.

from tree-math.