Change zipvmap to allow returning tuples of tensors

[benikm91]

This PR makes it possible to return tuple of tensors from zipvmap:

val (newXs, newYs) = zipvmap(Axis[Samples])(xs, ys): (x, y) =>
      (xs + ys, ys - xs) // return tuple

I specifically required this for my n-body-problem simulation. I think generally useful.

def step(masses: Tensor1[Particle, Float32], positions: Tensor2[Particle, Spatial, Float32], velocities: Tensor2[Particle, Spatial, Float32]): (Tensor1[Particle, Float32], Tensor2[Particle, Spatial, Float32], Tensor2[Particle, Spatial, Float32]) =
    zipvmap(Axis[Particle])(masses, positions, velocities): (m, p, v) =>
      val state = ParticleState(m, p, v)
      val newState = stepFor(state, masses, positions)
      (newState.mass, newState.position, newState.velocity)

(the example is in plotwit)

[marcelluethi]

I think the change make sense. But if we have it for zipvmap, we should also have it for vmap. Furthermore, I think the types should be made such that they are automatically deduced from the output type of the function (see comment below) I quickly generated a prototype of a solution using AI to test the feasibility.
The final signature looks something like:

    def zipvmap[L: Label, Inputs <: Tuple, FOut]( axis: Axis[L])(tensors: Inputs)
    (using ev: SharedAxisRemover[ShapesOf[Inputs], L])(
        f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut
    )(using   outEv: VmapResult[L, FOut])
    ): PrependAxis[L, FOut] =

Here PrependAxis would prepend the axis L to all elements of the result type.

[marcelluethi]

The Type MOut is a free Parameter. It must be possible to deduce it directly from the Output type of the function (using match types)

[benikm91]

Don't we do this in ZipVmapResult? For each case, e.g., for single Tensor its:

ZipVmapResult[L, Tensor[Shape, V], Tensor[L *: Shape, V]]

[marcelluethi]

Maybe, but there should not be a type parameter the user sees that has no obvious function. L, Inputs and FOut is really all that is needed to determine the output Type. These are, at least for the trained user, relatively straight-forward to interpret. In contrast to that, MOut is a parameter that seems unrelated to anything the user inputs and hence hard to understand. Furthermore, I suspect that this makes it also harder for type inference to work correctly.