def fprop_test()

in blocksparse/matmul.py [0:0]

• 22 lines of code
• 9 McCabe index (conditional complexity)

    def fprop_test(self, I, W, gate=None):
        bsize = self.bsize
        if self.axis:
            O = np.zeros((I.shape[0], self.KB, bsize))
            I = I.reshape((-1, self.CB, bsize))
            for k, lut in self.fprop_list:
                for c, w in lut:
                    O[:,k,:] += np.dot( I[:,c,:], W[w,:,:] ) # NC x CK = NK
            return O.reshape(I.shape[0], -1)
        else:
            N = I[0].size
            O = np.zeros((self.KB, bsize, N))
            I = I.reshape((self.CB, bsize, N))
            for k, lut in self.fprop_list:
                if gate is None:
                    for c, w in lut:
                        O[k,:,:] += np.dot( W[w,:,:].T, I[c,:,:] ) # CK.T x CN = KN
                else:
                    for c, w in lut:
                        if gate[w] != 0.0:
                            O[k,:,:] += np.dot( W[w,:,:].T, I[c,:,:] ) * gate[w] # CK.T x CN = KN

            return O.reshape(-1, N)