test_blas.py 文件源码

def test_dot22scalar():
    def cmp(a_shp, b_shp):
        a = tensor.fmatrix()
        b = tensor.fmatrix()
        scalar = tensor.fscalar()
        av = my_rand(*a_shp)
        bv = my_rand(*b_shp)

        f = theano.function(
            [a, b],
            tensor.dot(a, b) * numpy.asarray(4, 'float32'),
            mode=mode_with_gpu)
        f2 = theano.function(
            [a, b],
            tensor.dot(a, b) * numpy.asarray(4, 'float32'))
        t = f.maker.fgraph.toposort()
        assert any([isinstance(n.op, tcn.blas.GpuDot22Scalar) for n in t])
#        assert any([isinstance(n.op, tcn.basic_ops.GpuAllocEmpty)
#                    for n in t])
        assert numpy.allclose(f(av, bv), f2(av, bv))

        f = theano.function([a, b, scalar], tensor.dot(a, b) * scalar,
                            mode=mode_with_gpu)
        f2 = theano.function([a, b, scalar], tensor.dot(a, b) * scalar)
        t = f.maker.fgraph.toposort()
        assert any([isinstance(n.op, tcn.blas.GpuDot22Scalar) for n in t])
#        assert any([isinstance(n.op, tcn.basic_ops.GpuAllocEmpty)
#                    for n in t])
        assert numpy.allclose(f(av, bv, 0.5), f2(av, bv, 0.5))

        f = theano.function([a, b, scalar],
                            tensor.blas._dot22scalar(a, b, scalar),
                            mode=mode_with_gpu)
        f2 = theano.function([a, b, scalar], tensor.dot(a, b) * scalar)
        t = f.maker.fgraph.toposort()
        assert len(t) == 4
        assert isinstance(t[0].op, tcn.GpuFromHost)
        assert isinstance(t[1].op, tcn.GpuFromHost)
        assert isinstance(t[2].op, tcn.blas.GpuDot22Scalar)
        assert isinstance(t[3].op, tcn.HostFromGpu)
        assert numpy.allclose(f(av, bv, 0.5), f2(av, bv, 0.5))
    cmp((3, 4), (4, 5))
    cmp((0, 4), (4, 5))
    cmp((3, 4), (4, 0))
    cmp((3, 0), (0, 5))
    cmp((0, 4), (4, 0))
    cmp((0, 0), (0, 0))