INNER CODE UNIT · Python
block_fill_pytorch
microsoft/muzic · museformer/museformer/kernels/block_fill/main.py:90
def block_fill_pytorch(out, block_ranges, head_masks, fill_value):
num_blocks, num_heads = _check_input(block_ranges, head_masks, out)
assert out.shape[0] == num_heads
num_heads_2, _, _ = out.shape
assert num_heads_2 == num_heads
for (query_begin, query_end, key_begin, key_end), line_mask in zip(block_ranges, head_masks):
line_head_idx = torch.nonzero(line_mask, as_tuple=False).squeeze(-1) # (num,)
out[line_head_idx, query_begin: query_end, key_begin: key_end] = fill_value
return out