INNER CODE UNIT · Python

_check_input

microsoft/muzic · museformer/museformer/kernels/block_fill/main.py:55

def _check_input(block_ranges, head_masks, out):
    assert head_masks.dtype == torch.bool
    assert block_ranges.device == head_masks.device
    assert block_ranges.device == out.device
    num_blocks, col = block_ranges.shape
    num_blocks_2, num_heads = head_masks.shape
    assert num_blocks == num_blocks_2
    assert col == 4
    assert num_heads > 0
    return num_blocks, num_heads


def block_fill_cuda(out, block_ranges, head_masks, fill_value):
    num_blocks, num_heads = _check_input(block_ranges, head_masks, out)

    num_heads_2, seq_len_1, seq_len_2 = out.shape
    assert num_heads_2 == num_heads
    assert block_ranges.is_contiguous()

View source record →