INNER CODE UNIT · Python

range_fill_cuda

microsoft/muzic · museformer/museformer/kernels/range_fill/main.py:65

def range_fill_cuda(ranges, values, seq_len, pad_value, dtype=torch.long):
    if dtype not in (torch.long,):
        raise NotImplementedError
    num_ranges = ranges.shape[0]
    out = torch.full((seq_len,), pad_value, dtype=dtype, device=values.device)
    if num_ranges == 0:
        return out
    module = load_cuda_module()
    module.cuda_forward(ranges, values, num_ranges, out)
    return out


def range_fill_pytorch(ranges, values, seq_len, pad_value, dtype=torch.long):
    out = torch.full((seq_len,), pad_value, dtype=dtype, device=values.device)
    for idx, (begin, end) in enumerate(ranges):
        out[begin: end] = values[idx]
    return out

View source record →