def __init__()

in blocksparse/transformer.py [0:0]

• 57 lines of code
• 13 McCabe index (conditional complexity)

    def __init__(self, layout, block_size=64, heads=None, mask_callback=None, name=None):

        if len(layout.shape) == 2:
            assert heads is not None, "heads must be explicitly specified when using shared layouts per head"
            # broadcast same layout over all heads
            layout = np.expand_dims(layout, 0)

        if heads is None:
            heads = layout.shape[0]

        assert block_size in (8,16,32,64), "Block sizes of 8, 16, 32 and 64 currently supported"
        assert len(layout.shape) == 3, "bad layout shape: " + str(layout.shape)

        #self.layout       = layout > 0  # save boolean version for serialization purposes, TODO: save packbits or csr version
        self.blk_size     = block_size
        self.name         = name
        self.heads        = heads
        self.lut_heads    = layout.shape[0]
        self.ctx_blks_q   = layout.shape[1]
        self.ctx_blks_k   = layout.shape[2]
        self.blk_shape    = (block_size, block_size)
        self.nn_max       = 0
        self.tn_max       = 0
        self.softmax_dtype = None

        if layout.dtype != np.int32:
            layout = layout.astype(np.int32)

        self.nt_lut  = list()
        self.nn_lut  = list()
        self.tn_lut  = list()
        self.nt_list = list()
        self.nn_list = list()
        self.tn_list = list()
        blocks = None
        for head in range(layout.shape[0]):

            # convert to csr for vastly more efficient python iteration on large sparse layouts
            csr = sparse.csr_matrix(layout[head,:,:])
            ys, xs, bs = sparse.find(csr) # xs is in sorted order by default
            if blocks is None:
                blocks = len(bs)
            else:
                assert len(bs) == blocks, "number of layout blocks must be equal across heads"

            # make blocks contiguous along the rows (softmax code leverages this for increased performance)
            nt_list = sorted( zip(ys, xs) )
            ys = [b[0] for b in nt_list]
            xs = [b[1] for b in nt_list]

            nt_lut = np.array(nt_list, dtype=np.int32)
            nn_lut, nn_list, nn_max = self.xn_lut(ys, xs, blocks, self.ctx_blks_q)
            tn_lut, tn_list, tn_max = self.xn_lut(xs, ys, blocks, self.ctx_blks_k)

            self.nt_lut.append(nt_lut)
            self.nn_lut.append(nn_lut)
            self.tn_lut.append(tn_lut)
            self.nt_list.append(nt_list)
            self.nn_list.append(nn_list)
            self.tn_list.append(tn_list)
            self.nn_max = max(self.nn_max, nn_max)
            self.tn_max = max(self.tn_max, tn_max)

        self.blocks = blocks
        self.nt_lut = np.array(self.nt_lut, dtype=np.int32)
        self.nn_lut = np.array(self.nn_lut, dtype=np.int32)
        self.tn_lut = np.array(self.tn_lut, dtype=np.int32)

        if mask_callback is not None:
            self.init_softmax_mask(mask_callback)
        else:
            self.softmax_mask    = None
            self.softmax_mask_np = None