togethercomputer/ParallelKernelBench_Problems huggingface dataset:
# 46_gemv_decode
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    hidden_states: torch.Tensor,
    weight_shard: torch.Tensor,
    bias_shard: torch.Tensor,
) -> torch.Tensor:
    world_size = dist.get_world_size()

    local_logits = torch.matmul(hidden_states, weight_shard.t())
    local_logits = local_logits + bias_shard

    gathered = [torch.empty_like(local_logits) for _ in range(world_size)]
    dist.all_gather(gathered, local_logits.contiguous())
    logits = torch.cat(gathered, dim=1)

    return logits
# 47_fsdp_step_e2e
from __future__ import annotations

import math
from typing import Sequence

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors


def solution(
    X_local: Tensor,
    y_local: Tensor,
    flat_param_shard: Tensor,
    param_shapes: Sequence[tuple[int, ...]],
    exp_avg_shard: Tensor,
    exp_avg_sq_shard: Tensor,
    lr: float,
    beta1: float,
    beta2: float,
    eps: float,
    weight_decay: float,
    step: int,
) -> tuple[Tensor, Tensor, Tensor]:
    assert step >= 1

    world_size = dist.get_world_size()
    p = flat_param_shard.numel()

    device = flat_param_shard.device
    dtype = flat_param_shard.dtype

    templates = [torch.zeros(shape, dtype=dtype, device=device) for shape in param_shapes]
    full_flat = torch.empty(world_size * p, dtype=dtype, device=device)
    dist.all_gather_into_tensor(full_flat, flat_param_shard.contiguous())

    params_f = _unflatten_dense_tensors(full_flat, templates)
    params = [t.detach().requires_grad_(True) for t in params_f]

    h = F.relu(F.linear(X_local, params[0], params[1]))
    out = F.linear(h, params[2], params[3])
    loss = F.mse_loss(out, y_local)
    loss.backward()

    flat_g = _flatten_dense_tensors([x.grad for x in params])
    g_shard = torch.empty(p, dtype=flat_g.dtype, device=flat_g.device)
    dist.reduce_scatter_tensor(g_shard, flat_g.contiguous(), op=dist.ReduceOp.SUM)
    g_shard.div_(world_size)

    m = exp_avg_shard.clone()
    v = exp_avg_sq_shard.clone()
    theta = flat_param_shard.clone()

    m.mul_(beta1).add_(g_shard, alpha=1.0 - beta1)
    v.mul_(beta2).addcmul_(g_shard, g_shard, value=1.0 - beta2)
    bc1 = 1.0 - math.pow(beta1, step)
    bc2 = 1.0 - math.pow(beta2, step)
    m_hat = m / bc1
    v_hat = v / bc2
    denom = v_hat.sqrt().add(eps)

    theta.add_(m_hat.div(denom), alpha=-lr)
    theta.add_(flat_param_shard, alpha=-lr * weight_decay)

    return theta, m, v
# 48_fsdp_and_tp
from __future__ import annotations

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor


def _make_tp_fsdp_groups(n_tp: int, n_fsdp: int, rank: int):
    tp_group = None
    fsdp_group = None
    for j in range(n_fsdp):
        ranks = [j * n_tp + ii for ii in range(n_tp)]
        g = dist.new_group(ranks)
        if rank in ranks:
            tp_group = g
    for i in range(n_tp):
        ranks = [jj * n_tp + i for jj in range(n_fsdp)]
        g = dist.new_group(ranks)
        if rank in ranks:
            fsdp_group = g
    assert tp_group is not None and fsdp_group is not None
    return tp_group, fsdp_group


def _gather_fsdp_concat_dim0(shard: Tensor, fsdp_group, parts: int) -> Tensor:
    lst = [torch.empty_like(shard) for _ in range(parts)]
    dist.all_gather(lst, shard.contiguous(), group=fsdp_group)
    return torch.cat(lst, dim=0)


def _gather_fsdp_concat_dim1(shard: Tensor, fsdp_group, parts: int) -> Tensor:
    lst = [torch.empty_like(shard) for _ in range(parts)]
    dist.all_gather(lst, shard.contiguous(), group=fsdp_group)
    return torch.cat(lst, dim=1)


@torch.no_grad()
def solution(
    x_local: Tensor,
    W1_shard: Tensor,
    W2_shard: Tensor,
    W3_shard: Tensor,
    n_tp: int,
    n_fsdp: int,
) -> Tensor:
    world_size = dist.get_world_size()
    rank = dist.get_rank()
    assert world_size == n_tp * n_fsdp, (
        f"world_size ({world_size}) must equal n_tp * n_fsdp ({n_tp} * {n_fsdp})"
    )

    tp_group, fsdp_group = _make_tp_fsdp_groups(n_tp, n_fsdp, rank)

    W1 = _gather_fsdp_concat_dim0(W1_shard, fsdp_group, n_fsdp)
    W2 = _gather_fsdp_concat_dim0(W2_shard, fsdp_group, n_fsdp)
    W3 = _gather_fsdp_concat_dim1(W3_shard, fsdp_group, n_fsdp)

    x1 = x_local @ W1
    x2 = x_local @ W2
    z = F.silu(x1) * x2
    y_partial = z @ W3

    y = y_partial.clone()
    dist.all_reduce(y, op=dist.ReduceOp.SUM, group=tp_group)
    return y
# 52_fp8_reduce_scatter_grads
from __future__ import annotations

import torch
import torch.distributed as dist
from torch import Tensor

_FP8_E4M3_MAX = 448.0


@torch.no_grad()
def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor:
    out = torch.roll(amax_history, shifts=-1, dims=0)
    out[-1] = cur_abs_max.to(dtype=out.dtype)
    return out


@torch.no_grad()
def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor:
    xf = x.float()
    qs = xf / scale
    q = qs.to(torch.float8_e4m3fn)
    return (q.float() * scale).to(dtype=x.dtype)


@torch.no_grad()
def solution(flat_grads: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]:
    world_size = dist.get_world_size()
    n = flat_grads.numel()
    shard_elems = n // world_size

    cur_abs_max = flat_grads.abs().max().to(torch.float32)
    updated_hist = _update_amax_history(amax_history, cur_abs_max)

    scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX
    recon = _fp8_round_trip_bf16(flat_grads, scale)

    out_shard = torch.empty(shard_elems, dtype=flat_grads.dtype, device=flat_grads.device)
    dist.reduce_scatter_tensor(out_shard, recon.contiguous(), op=dist.ReduceOp.SUM)
    out_shard.div_(world_size)

    return out_shard, updated_hist
# 53_fp8_allgather_params
from __future__ import annotations

import torch
import torch.distributed as dist
from torch import Tensor

_FP8_E4M3_MAX = 448.0


@torch.no_grad()
def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor:
    out = torch.roll(amax_history, shifts=-1, dims=0)
    out[-1] = cur_abs_max.to(dtype=out.dtype)
    return out


@torch.no_grad()
def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor:
    xf = x.float()
    qs = xf / scale
    q = qs.to(torch.float8_e4m3fn)
    return (q.float() * scale).to(dtype=x.dtype)


@torch.no_grad()
def solution(flat_param_shard: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]:
    world_size = dist.get_world_size()
    p = flat_param_shard.numel()

    cur_abs_max = flat_param_shard.abs().max().to(torch.float32)
    updated_hist = _update_amax_history(amax_history, cur_abs_max)

    scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX
    recon = _fp8_round_trip_bf16(flat_param_shard, scale)

    full = torch.empty(world_size * p, dtype=flat_param_shard.dtype, device=flat_param_shard.device)
    dist.all_gather_into_tensor(full, recon.contiguous())

    return full, updated_hist
# 58_openclip_contrastive_loss
from typing import Optional

import torch
import torch.distributed as dist
import torch.nn.functional as F


def _siglip_loss(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    logit_scale: float,
    logit_bias: float,
    negative_only: bool,
) -> torch.Tensor:
    batch = image_features.size(0)
    logits = logit_scale * image_features @ text_features.T + logit_bias

    if negative_only:
        return -F.logsigmoid(-logits).sum() / batch

    labels = -torch.ones((batch, batch), device=logits.device, dtype=logits.dtype)
    labels.diagonal().fill_(1)
    return -F.logsigmoid(labels * logits).sum() / batch


def _exchange(
    group: dist.ProcessGroup,
    recv_from: int,
    send_to: int,
    tensor: torch.Tensor,
) -> torch.Tensor:
    out = torch.empty_like(tensor)
    ops = [
        dist.P2POp(dist.isend, tensor, dist.get_global_rank(group, send_to), group=group),
        dist.P2POp(dist.irecv, out, dist.get_global_rank(group, recv_from), group=group),
    ]
    for req in dist.batch_isend_irecv(ops):
        req.wait()
    return out


def _exchange_bidir(
    group: dist.ProcessGroup,
    left: int,
    right: int,
    tensor_to_left: torch.Tensor,
    tensor_to_right: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    from_left = torch.empty_like(tensor_to_right)
    from_right = torch.empty_like(tensor_to_left)
    left_global = dist.get_global_rank(group, left)
    right_global = dist.get_global_rank(group, right)
    ops = [
        dist.P2POp(dist.isend, tensor_to_right, right_global, group=group),
        dist.P2POp(dist.isend, tensor_to_left, left_global, group=group),
        dist.P2POp(dist.irecv, from_right, right_global, group=group),
        dist.P2POp(dist.irecv, from_left, left_global, group=group),
    ]
    for req in dist.batch_isend_irecv(ops):
        req.wait()
    return from_right, from_left


@torch.no_grad()
def solution(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    logit_scale: float,
    logit_bias: float = 0.0,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    loss = _siglip_loss(image_features, text_features, logit_scale, logit_bias, False)

    left = (rank - 1) % world_size
    right = (rank + 1) % world_size
    text_to_left = text_features
    text_to_right = text_features
    num_bidir, remainder = divmod(world_size - 1, 2)

    for _ in range(num_bidir):
        from_right, from_left = _exchange_bidir(group, left, right, text_to_left, text_to_right)
        loss = loss + _siglip_loss(image_features, from_right, logit_scale, logit_bias, True)
        loss = loss + _siglip_loss(image_features, from_left, logit_scale, logit_bias, True)
        text_to_left, text_to_right = from_right, from_left

    if remainder:
        text_recv = _exchange(group, left, right, text_to_right)
        loss = loss + _siglip_loss(image_features, text_recv, logit_scale, logit_bias, True)

    return loss
# 59_physicsnemo_distributed_rfft
from typing import Optional, Sequence

import torch
import torch.distributed as dist


def _all_to_all_transpose(
    tensor: torch.Tensor,
    split_dim: int,
    group: dist.ProcessGroup,
) -> list[torch.Tensor]:
    world_size = dist.get_world_size(group)
    chunk = tensor.shape[split_dim] // world_size
    send = [x.contiguous() for x in torch.split(tensor, chunk, dim=split_dim)]
    recv = [torch.empty_like(send[0]) for _ in range(world_size)]
    dist.all_to_all(recv, send, group=group)
    return recv


def _truncate(tensor: torch.Tensor, dim: int, size: int) -> torch.Tensor:
    slices = [slice(None)] * tensor.ndim
    slices[dim % tensor.ndim] = slice(0, size)
    return tensor[tuple(slices)].contiguous()


@torch.no_grad()
def solution(
    x: torch.Tensor,
    s: Sequence[int],
    dim: Sequence[int],
    norm: str = "ortho",
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    dim0, dim1 = int(dim[0]), int(dim[1])

    n0 = int(s[0]) if s[0] is not None else None
    n1 = int(s[1]) if s[1] is not None else None

    x1 = torch.fft.fft(x, n=n0, dim=dim0, norm=norm)

    x1_recv = _all_to_all_transpose(x1, split_dim=dim0, group=group)
    x1_tran = torch.cat(x1_recv, dim=dim1)

    x2 = torch.fft.fft(x1_tran, n=n1, dim=dim1, norm=norm)

    return _truncate(x2, dim1, x2.shape[dim1] // 2 + 1)
# 63_deepmd_kalman_filter_optimizer
from typing import List, Tuple

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    H: List[torch.Tensor],
    error: torch.Tensor,
    weights: List[torch.Tensor],
    P: List[torch.Tensor],
    kalman_lambda: float,
    kalman_nue: float = 0.9987,
) -> Tuple[List[torch.Tensor], List[torch.Tensor], torch.Tensor]:
    weights_num = len(weights)
    lam = torch.as_tensor(kalman_lambda, dtype=weights[0].dtype, device=weights[0].device)
    err = error.to(device=weights[0].device, dtype=weights[0].dtype)

    tmp = 0
    for i in range(weights_num):
        tmp = tmp + (lam + torch.matmul(torch.matmul(H[i].T, P[i]), H[i]))

    if dist.is_initialized():
        dist.all_reduce(tmp, op=dist.ReduceOp.SUM)

    A = 1 / tmp

    for i in range(weights_num):
        K = torch.matmul(P[i], H[i])
        weights[i] = weights[i] + A * err * K
        P[i] = (1 / lam) * (P[i] - A * torch.matmul(K, K.T))

    if dist.is_initialized():
        device = weights[0].device
        local_shape = [tensor.shape[0] for tensor in weights]

        shape_list = [
            torch.zeros_like(torch.empty(1), dtype=torch.float64, device=device)
            for _ in range(dist.get_world_size())
        ]
        dist.all_gather_object(shape_list, local_shape)

        weight_tensor = torch.cat([w.reshape(-1) for w in weights], dim=0).to(torch.float64)
        world_shape = [sum(inner_list) for inner_list in shape_list]

        weight_list = [
            torch.zeros(world_shape[i], dtype=torch.float64, device=device)
            for i in range(len(world_shape))
        ]
        dist.all_gather(weight_list, weight_tensor)

        result: List[torch.Tensor] = []
        for i in range(dist.get_world_size()):
            result = result + [
                t.reshape(-1, 1).to(dtype=weights[0].dtype)
                for t in torch.split(weight_list[i], shape_list[i])
            ]
        weights = result

    kalman_lambda_next = (
        torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device) * lam
        + 1
        - torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device)
    )

    return weights, P, kalman_lambda_next
# 65_gnn_feature_exchange_all2all
from typing import List, Optional

import torch
import torch.distributed as dist


def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:
    cutoff = len(chunks) - dist.get_rank(group)
    return chunks[cutoff:] + chunks[:cutoff]


def _all_to_all(
    outputs: List[torch.Tensor],
    inputs: List[torch.Tensor],
    group: dist.ProcessGroup,
) -> None:
    outputs = _shift(list(outputs), group)
    inputs = _shift(list(inputs), group)
    if outputs and outputs[0].is_cuda:
        dist.all_to_all(outputs, inputs, group=group)
        return

    output_splits = [out.size(0) for out in outputs]
    input_splits = [inp.size(0) for inp in inputs]
    flat_out = torch.cat(outputs) if outputs else torch.empty(0)
    flat_in = torch.cat(inputs) if inputs else torch.empty(0)
    dist.all_to_all_single(
        flat_out,
        flat_in,
        output_split_sizes=output_splits,
        input_split_sizes=input_splits,
        group=group,
    )
    for out, temp in zip(outputs, flat_out.split(output_splits)):
        out.copy_(temp)


@torch.no_grad()
def solution(
    local_features: torch.Tensor,
    seed_inverse_ids: torch.Tensor,
    counts_sent: List[int],
    counts_received: List[int],
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    gathered = local_features[seed_inverse_ids]
    out = local_features.new_empty((sum(counts_sent),) + local_features.shape[1:])
    _all_to_all(
        list(torch.split(out, counts_sent)),
        list(torch.split(gathered, counts_received)),
        group,
    )
    return out
# 66_gnn_feature_exchange_all2all_backward
from typing import List, Optional

import torch
import torch.distributed as dist


def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:
    cutoff = len(chunks) - dist.get_rank(group)
    return chunks[cutoff:] + chunks[:cutoff]


def _all_to_all(
    outputs: List[torch.Tensor],
    inputs: List[torch.Tensor],
    group: dist.ProcessGroup,
) -> None:
    outputs = _shift(list(outputs), group)
    inputs = _shift(list(inputs), group)
    if outputs and outputs[0].is_cuda:
        dist.all_to_all(outputs, inputs, group=group)
        return

    output_splits = [out.size(0) for out in outputs]
    input_splits = [inp.size(0) for inp in inputs]
    flat_out = torch.cat(outputs) if outputs else torch.empty(0)
    flat_in = torch.cat(inputs) if inputs else torch.empty(0)
    dist.all_to_all_single(
        flat_out,
        flat_in,
        output_split_sizes=output_splits,
        input_split_sizes=input_splits,
        group=group,
    )
    for out, temp in zip(outputs, flat_out.split(output_splits)):
        out.copy_(temp)


@torch.no_grad()
def solution(
    grad_output: torch.Tensor,
    seed_inverse_ids: torch.Tensor,
    seed_size: int,
    counts_sent: List[int],
    counts_received: List[int],
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    out = grad_output.new_empty((sum(counts_received),) + grad_output.shape[1:])
    _all_to_all(
        list(torch.split(out, counts_received)),
        list(torch.split(grad_output, counts_sent)),
        group,
    )

    idx = torch.empty((2, out.shape[0]), dtype=torch.int64, device=grad_output.device)
    idx[0] = seed_inverse_ids
    idx[1] = torch.arange(out.shape[0], device=grad_output.device)
    coo = torch.sparse_coo_tensor(
        idx,
        torch.ones(idx.shape[1], dtype=grad_output.dtype, device=idx.device),
        size=(seed_size, idx.shape[1]),
    )
    return torch.sparse.mm(coo, out)
# 67_gnn_sparse_embedding_all2all
from typing import Optional, Tuple

import torch
import torch.distributed as dist


def _generate_permutation_remainder(
    idx: torch.Tensor,
    world_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    owner = (idx % world_size).long()
    send_splits = torch.bincount(owner, minlength=world_size)
    perm = torch.argsort(owner, stable=True).long()
    return perm, send_splits


@torch.no_grad()
def solution(
    idx: torch.Tensor,
    value: torch.Tensor,
    num_nodes: int,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return idx, value

    perm, send_splits = _generate_permutation_remainder(idx, world_size)

    recv_splits = torch.empty_like(send_splits)
    dist.all_to_all_single(recv_splits, send_splits, group=group)

    recv_splits = recv_splits.to("cpu", non_blocking=True)
    send_splits = send_splits.to("cpu", non_blocking=True)
    send_idx = idx[perm]
    send_value = value[perm]
    if idx.is_cuda:
        torch.cuda.current_stream().synchronize()

    recv_count = int(recv_splits.sum().item())
    recv_splits_list = recv_splits.tolist()
    send_splits_list = send_splits.tolist()

    recv_idx = torch.empty((recv_count,), dtype=idx.dtype, device=idx.device)
    dist.all_to_all_single(
        recv_idx,
        send_idx,
        output_split_sizes=recv_splits_list,
        input_split_sizes=send_splits_list,
        group=group,
    )

    recv_value = torch.empty(
        (recv_count, *value.shape[1:]),
        dtype=value.dtype,
        device=value.device,
    )
    dist.all_to_all_single(
        recv_value,
        send_value,
        output_split_sizes=recv_splits_list,
        input_split_sizes=send_splits_list,
        group=group,
    )

    return recv_idx, recv_value
# 68_gnn_sparse_feature_fetch_projection
from typing import Optional

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    local_embedding_shard: torch.Tensor,
    input_node_ids: torch.Tensor,
    proj_matrix: torch.Tensor,
    num_total_nodes: int,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    rank = dist.get_rank(group)
    shard_size = (num_total_nodes + world_size - 1) // world_size
    embed_dim = local_embedding_shard.shape[1]
    num_queries = input_node_ids.shape[0]

    owner = (input_node_ids // shard_size).clamp(max=world_size - 1)
    sort_idx = torch.argsort(owner, stable=True)
    sorted_ids = input_node_ids[sort_idx]
    sorted_owner = owner[sort_idx]

    send_counts = torch.zeros(
        world_size, dtype=torch.long, device=input_node_ids.device
    )
    send_counts.scatter_add_(0, sorted_owner, torch.ones_like(sorted_owner))
    recv_counts = torch.empty_like(send_counts)
    dist.all_to_all_single(recv_counts, send_counts, group=group)

    send_splits = send_counts.to("cpu").tolist()
    recv_splits = recv_counts.to("cpu").tolist()
    recv_ids = torch.empty(
        int(recv_counts.sum().item()),
        dtype=input_node_ids.dtype,
        device=input_node_ids.device,
    )
    dist.all_to_all_single(
        recv_ids,
        sorted_ids,
        output_split_sizes=recv_splits,
        input_split_sizes=send_splits,
        group=group,
    )

    local_ids = (recv_ids - rank * shard_size).long()
    fetched = local_embedding_shard[local_ids]

    gathered = torch.empty(
        (num_queries, embed_dim), dtype=fetched.dtype, device=fetched.device
    )
    dist.all_to_all_single(
        gathered,
        fetched,
        output_split_sizes=send_splits,
        input_split_sizes=recv_splits,
        group=group,
    )
    emb = gathered[torch.argsort(sort_idx, stable=True)]

    return emb @ proj_matrix
# 69_gnn_negative_scoring
from typing import Optional

import torch
import torch.distributed as dist


def _broadcast_data(
    rank: int,
    world_size: int,
    data: torch.Tensor,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    if world_size == 1:
        return data

    sizes = torch.zeros(world_size, dtype=torch.long, device=data.device)
    sizes[rank] = data.shape[0]
    dist.all_reduce(sizes, op=dist.ReduceOp.SUM, group=group)

    send_splits = [data.shape[0] for _ in range(world_size)]
    recv_splits = sizes.to("cpu").tolist()
    send = data.repeat(*([world_size] + [1] * (data.ndim - 1))).contiguous()
    recv = torch.empty(
        (int(sizes.sum().item()), *data.shape[1:]),
        dtype=data.dtype,
        device=data.device,
    )
    dist.all_to_all_single(
        recv,
        send,
        output_split_sizes=recv_splits,
        input_split_sizes=send_splits,
        group=group,
    )
    return recv


def _calc_ranking(pos_score: torch.Tensor, neg_score: torch.Tensor) -> torch.Tensor:
    scores = torch.cat([pos_score.view(-1, 1), neg_score], dim=1)
    _, indices = torch.sort(torch.sigmoid(scores), dim=1, descending=True)
    return torch.nonzero(indices == 0)[:, 1].view(-1).detach() + 1


@torch.no_grad()
def solution(
    local_pos_scores: torch.Tensor,
    local_neg_scores: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    pos_scores = _broadcast_data(rank, world_size, local_pos_scores, group)
    neg_scores = _broadcast_data(rank, world_size, local_neg_scores, group)
    return _calc_ranking(pos_scores, neg_scores)