def METHOD_NAME(ctx, grad_output):
    result = (
        torch.empty_like(grad_output)
        if ctx.input_splits is None
        else grad_output.new_empty(
            size=[sum(ctx.input_splits)] + list(grad_output.size()[1:])
        )
    )
    torch.distributed.all_to_all_single(
        result,
        grad_output,
        output_split_sizes=ctx.input_splits,
        input_split_sizes=ctx.output_splits,
    )
    return result, None, None