NamedTuples failing for zipvmap

[benikm91]

The current zipvmap and vmap do not yet support named tuples as the return type of f. I failed to achieve this by extending the PrependAxis match type. This PR provides a solution, but I am unsure whether removing the match type is necessary. This PR is a conversation starter about whether and how we can add named tuple support more cleanly with match types. And highlighting the lack of support.

Named tuples could be quite powerful, as the tuple names can be tracked across vmap and zipvmap operations. This is not possible for case classes with a fixed tensor structure. Here with explicit type alias to clarify what is (automatically) happening:

val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f)
type FOut = (first: Tensor1[B, Float32], second: Tensor1[B, Float32]) // return named tuple of tensor
def f(x: Tensor1[B, Float32]): FOut =
  (first = x +! 5f, second = x -! 5f)
type TrackedOut = (first: Tensor2[A, B, Float32], second: Tensor2[A, B, Float32]) // automatically track names across tensor expansion
val res: TrackedOut = t.vmap(Axis[A]): x => // vmap automatically transforms FOut to TrackedOut (adds A axis to each field) keeps names
  f(x)

[marcelluethi]

I wonder if there is a disadvantage of given up match types in terms of type inference. Have you noticed anything?
In case type inference or compile time messages are getting worse, I think it is not worth the trouble. It seems to me that in 99% of the cases, what is returned is just a single value, not a tuple. Now we support tuples of all shapes. And even if we could support named tuples, there is still the gap that we cannot support case classes.

We could always work around it by manually calling to tuple

val res = zipvmap(Axis[A])(params.toTuple) { case (x1, x2) => x1 + x2 }
val (t1, t2) = t.vmap(Axis[A])(x => (x +! 5f, x -! 5f))
val named = (name1= t1,  name2 = t2)

However, if cost is only implementation complexity, I am all in for it. What is your experience?