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]]