submission 512298

#!POPCORN leaderboard prefixsum_v2
#!POPCORN gpu A100

import torch
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of inclusive prefix sum using PyTorch.
    Args:
        data: Input tensor to compute prefix sum on
    Returns:
        Tensor containing the inclusive prefix sum
    """
    data, output = data
    output[...] = torch.cumsum(data, dim=0)
    return output