zigzag_gather_from_group_ranks

zigzag_gather_from_group_ranks(data, group, seq_dim=0)

Reconstructs complete tensor from zigzag-distributed chunks.

Takes data distributed across ranks in zigzag pattern and reassembles the original complete tensor.

Parameters:

• data – tensor fragment from current rank to be gathered.

• group – the group to gather data from.

• seq_dim – dimension along which to concatenate fragments.

Returns:

Reconstructed tensor with fragments from all ranks.