test_constant.py 文件源码

def test_cputensor_add():
    """TODO."""
    Y = ng.make_axis(length=2)
    M = ng.make_axis(length=2)
    N = ng.make_axis(length=2)

    a = ng.constant(np.array([3, 5], dtype=np.float32), [Y])
    b = ng.constant(np.array([3, 5], dtype=np.float32), [Y])
    c = a + b
    with executor(c) as ex:
        result = ex()
    assert np.array_equal(result, [6, 10])

    np_a = np.array([[1, 2], [3, 4]], dtype=np.float32)
    np_b = np.array([[1, 2], [3, 4]], dtype=np.float32)
    np_c = np_a + np_b

    a = ng.constant(np_a, [M, N])
    b = ng.constant(np_b, [M, N])
    c = a + b
    with executor(c) as ex:
        result = ex()
    assert np.array_equal(result, np_c)