INNER CODE UNIT · Python
max_query
microsoft/muzic · museformer/museformer/kernels/block_fill/main.py:76
max_query = (block_ranges[:, 1] - block_ranges[:, 0]).max().item()
max_key = (block_ranges[:, 3] - block_ranges[:, 2]).max().item()
if max_query <= 0 or max_key <= 0:
return out
module = load_cuda_module()
module.cuda_forward(block_ranges, head_masks, num_blocks, max_query, max_key, num_heads,
seq_len_1, seq_len_2, fill_value, out)
return out
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