INNER CODE UNIT · Python
segment_arange_cuda
microsoft/muzic · museformer/museformer/kernels/segment_arange/main.py:64
def segment_arange_cuda(ranges, start, 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=ranges.device)
if num_ranges == 0:
return out
module = load_cuda_module()
module.cuda_forward(ranges, start, num_ranges, out)
return out
def segment_arange_pytorch(ranges, start, seq_len, pad_value, dtype=torch.long):
out = torch.full((seq_len,), pad_value, dtype=dtype, device=ranges.device)
for idx, (begin, end) in enumerate(ranges):
out[begin: end] = torch.arange(start, start + (end - begin), dtype=dtype, device=ranges.device)
return out