Fabric.all_gather should support concatenation

Description & Motivation

When using DistributedDataParallel, the "Fabric.all_gather" option is useful to collect tensors from all devices. x = fabric.all_gather(x) will return a Tensor of dimensions [fabric.world_size , *input.shape], which is natural for a lot of further processing.

Pitch

It would be useful, if there was a flag fabric.all_gather(x, op=OP, dim=1) where OP can be "stack" and "cat". For example, it would be possible to directly output [batch_size_per_gpu * fabric_world_size, ...], which is more natural for processing the input to compute metrics etc. I assume it is also faster to directly concatenate instead of stacking first and then re-ordering.

Alternatives

Alterantive is to keep doing this

            if STRATEGY == "ddp":
                y_hat = fabric.all_gather(y_hat)
                y = fabric.all_gather(y)
                y_hat, y = y_hat.flatten(0,1), y.flatten(0,1)

Additional context

fabric.all_gather is not required for dp as far as I can tell, but it also doesn't hurt. Having to flatten the tensor, makes "ddp" code incompatible with "dp" code. Doing `fabric.all_gather(y_hat, "cat", dim=0) would work for both equally.

cc @borda @carmocca @justusschock @awaelchli

Start at the Fabric.all_gather entry point and trace how gathering behaves under DistributedDataParallel and data parallel strategies. Define the requested stack and concatenation behavior, including the dimension argument, then verify that the resulting tensor shapes support the metric-processing example and both strategies.