def bst_conv_layout()

in blocksparse/utils.py [0:0]

• 47 lines of code
• 15 McCabe index (conditional complexity)

def bst_conv_layout(input_h=1, input_w=1, filter_h=1, filter_w=1, stride=1, blk_size=32, autoregressive=True):

    H = input_h
    W = input_w
    R = filter_h
    S = filter_w

    assert H % stride == 0 or H == 1
    assert W % stride == 0

    P = H // stride or 1
    Q = W // stride

    if H == 1:
        R = 1
        pad_r = 0
    else:
        pad_r = -1
        for r in range(R):
            if P == out_dim(R, H, r, stride):
                pad_r = r
                break
    assert pad_r >= 0, "Even size filters only work with stride 2."

    pad_s = -1
    for s in range(S):
        if Q == out_dim(S, W, s, stride):
            pad_s = s
            break
    assert pad_s >= 0, "Even size filters only work with stride 2."

    print(f"P:{P} Q:{Q} H:{H} W:{W} R:{R} S:{S} std:{stride} pad_r:{pad_r} pad_s:{pad_s}")

    assert P*Q % blk_size == 0, f"P:{P} Q:{Q}"
    assert H*W % blk_size == 0, f"H:{H} W:{W}"

    mask_set = set()
    layout = np.zeros((P*Q//blk_size, H*W//blk_size), dtype=np.bool)

    # just compute the output pixels within the tile
    for p, q in np.ndindex(P, Q):
        for     h in conv_slice(p, H, R, pad_r, stride):
            for w in conv_slice(q, W, S, pad_s, stride):
                x = h*W + w
                y = p*Q + q
                if not autoregressive or p*stride*Q*stride + q*stride >= x:
                    layout[y//blk_size, x//blk_size] = 1
                    mask_set.add((y, x))

    def cb(blk_shape, head_idx, qry_idx, key_idx, blk_idx):
        mask = np.zeros(blk_shape, dtype=np.bool)

        q0 = qry_idx*blk_shape[0]
        k0 = key_idx*blk_shape[1]
        for q, k in np.ndindex(blk_shape):
            if (q0 + q, k0 + k) in mask_set:
                mask[q, k] = 1
        return mask

    return layout, cb