test_extra_ops.py 文件源码

def test_op(self):
        for axis, cond, shape in zip(self.axis_list, self.cond_list,
                                     self.shape_list):
            cond_var = theano.tensor.ivector()
            data = numpy.random.random(size=shape).astype(theano.config.floatX)
            data_var = theano.tensor.matrix()

            f = theano.function([cond_var, data_var],
                                self.op(cond_var, data_var, axis=axis))

            expected = numpy.compress(cond, data, axis=axis)
            tested = f(cond, data)

            assert tested.shape == expected.shape
            assert numpy.allclose(tested, expected)