def _build_edge_lut()

in blocksparse/conv.py [0:0]

• 35 lines of code
• 12 McCabe index (conditional complexity)

    def _build_edge_lut(self, MPQ, mpqLut):

        # Hash the mpq coordinates on unique edge overlap patterns
        # The hash key is the list of lut indicies where the offset is -1
        PQ = MPQ[1] * MPQ[2]
        Q  = MPQ[2]
        edge_map = dict()
        mLut, pLut, qLut = mpqLut
        for m,p,q in np.ndindex(*MPQ):
            key = list()
            for di, d in enumerate(mLut[m]):
                for hi, h in enumerate(pLut[p]):
                    for wi, w in enumerate(qLut[q]):
                        if any(x == -1 for x in (d,h,w)):
                            key.append((di,hi,wi))
            if len(key):
                key = tuple(key)
                mpqOffset = m*PQ + p*Q + q
                edge_list = edge_map.get(key)
                if edge_list is None:
                    edge_map[key] = [mpqOffset]
                else:
                    edge_list.append(mpqOffset)

        self.edgeBiasDim = len(edge_map)

        if self.edgeBiasDim:
            # so K x len(edge_map) is the size of the bias vector
            # we need a lut of bias index => mpqOffset mappings
            biasHead = list()
            biasData = list()
            biasMap  = sorted(edge_map.values(), key=lambda x: x[0])
            offset   = len(biasMap) * 2
            # the lut contains a header with 2 entries per unique bias: offset, size
            for mpqList in biasMap:
                biasHead.extend((offset, len(mpqList)))
                biasData.extend(mpqList)
                offset += len(mpqList)

            pad4 = 4 - (len(biasData) & 3) if (len(biasData) & 3) else 0
            biasLut = biasHead + biasData + ( [0] * pad4 )
            self.edgeEntries = len(biasData)
            self.edgeBiasMap = biasMap
            self.edgeBiasLut = tf.constant(np.array(biasLut, dtype=np.int32), name="edge_bias_lut")