def all_to_all(outputs, inputs, group=None, async_op=False):
    """Wrapper for thd.all_to_all that permuted outputs and inputs before
    calling it. The arguments have the permutation
    `rank, ..., world_size - 1, 0, ..., rank - 1` and we make it
    `0, world_size - 1` before calling `thd.all_to_all`."""
    shift_fn = partial(_shift, group=group)
    outputs = shift_fn(list(outputs))
    inputs = shift_fn(list(inputs))
    if outputs[0].is_cuda:
        return thd.all_to_all(outputs, inputs, group, async_op)
    # gloo backend will be used.
    outputs_single = torch.cat(outputs)
    output_split_sizes = [o.size(0) for o in outputs]
    handle = thd.all_to_all_single(
        outputs_single,
        torch.cat(inputs),
        output_split_sizes,
        [i.size(0) for i in inputs],
        group,
        async_op,
    )
    temp_outputs = outputs_single.split(output_split_sizes)

    class _Waiter:
        def __init__(self, handle, outputs, temp_outputs):
            self.handle = handle
            self.outputs = outputs
            self.temp_outputs = temp_outputs

        def wait(self):
            """Returns the stored value when invoked."""
            handle = self.handle
            outputs = self.outputs
            temp_outputs = self.temp_outputs
            # Ensure that there is no leak
            self.handle = self.outputs = self.temp_outputs = None

            if handle is not None:
                handle.wait()
            for output, temp_output in zip(outputs, temp_outputs):
                output.copy_(temp_output)

    post_processor = _Waiter(handle, outputs, temp_outputs)
    return post_processor if async_op else post_processor.wait()
Wrapper for thd.all_to_all that permuted outputs and inputs before
    calling it. The arguments have the permutation
    `rank, ..., world_size - 1, 0, ..., rank - 1` and we make it
    `0, world_size - 1` before calling `thd.all_to_all`.
python/dgl/graphbolt/subgraph_sampler.py