INNER CODE UNIT · Python
block_fill_
microsoft/muzic · museformer/museformer/kernels/block_fill/main.py:35
def block_fill_(out, block_ranges, head_masks, fill_value, no_cuda_kernel=False):
"""
:param block_ranges: (num_blocks, 4)
:param head_masks: (num_blocks, num_heads)
:param out:
:param fill_value:
:param no_cuda_kernel:
:return:
"""
if block_ranges.shape[0] == 0:
return out
assert isinstance(fill_value, bool)
if block_ranges.is_cuda and not no_cuda_kernel:
return block_fill_cuda(out, block_ranges, head_masks, fill_value)
else:
return block_fill_pytorch(out, block_ranges, head_masks, fill_value)