示例#1
0
def test_codegen_array_contraction_construction():
    cg = CodegenArrayContraction(A)
    assert cg == A

    s = Sum(A[i] * B[i], (i, 0, 3))
    cg = parse_indexed_expression(s)
    assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B),
                                         (0, 1))

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (1, 0))
    assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B),
                                         (0, 1))

    expr = M * N
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
    assert CodegenArrayContraction.from_MatMul(expr) == result
    elem = expr[i, j]
    assert parse_indexed_expression(elem) == result

    expr = M * N * M
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, M),
                                     (1, 2), (3, 4))
    assert CodegenArrayContraction.from_MatMul(expr) == result
    elem = expr[i, j]
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N),
                                     (1, 4), (2, 5))
    cg = parse_indexed_expression(elem)
    cg = cg.sort_args_by_name()
    assert cg == result
示例#2
0
def test_codegen_array_recognize_matrix_mul_lines():

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
    assert recognize_matrix_expression(cg) == Trace(M)

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
                                 (2, 3))
    assert recognize_matrix_expression(cg) == [Trace(M), Trace(N)]

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
                                 (1, 2))
    assert recognize_matrix_expression(cg) == Trace(M * N)

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
                                 (1, 3))
    assert recognize_matrix_expression(cg) == Trace(M * N.T)

    cg = parse_indexed_expression((M * N * P)[i, j])
    assert recognize_matrix_expression(cg) == M * N * P
    cg = CodegenArrayContraction.from_MatMul(M * N * P)
    assert recognize_matrix_expression(cg) == M * N * P

    cg = parse_indexed_expression((M * N.T * P)[i, j])
    assert recognize_matrix_expression(cg) == M * N.T * P
    cg = CodegenArrayContraction.from_MatMul(M * N.T * P)
    assert recognize_matrix_expression(cg) == M * N.T * P

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
                                 (5, 6))
    assert recognize_matrix_expression(cg) == [M * N, P * Q]

    expr = -2 * M * N
    elem = expr[i, j]
    cg = parse_indexed_expression(elem)
    assert recognize_matrix_expression(cg) == -2 * M * N
示例#3
0
def test_codegen_array_recognize_matrix_mul_lines():

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
    assert recognize_matrix_expression(cg) == Trace(M)

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1), (2, 3))
    assert recognize_matrix_expression(cg) == [Trace(M), Trace(N)]

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3), (1, 2))
    assert recognize_matrix_expression(cg) == Trace(M*N)

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3))
    assert recognize_matrix_expression(cg) == Trace(M*N.T)

    cg = parse_indexed_expression((M*N*P)[i,j])
    assert recognize_matrix_expression(cg) == M*N*P
    cg = CodegenArrayContraction.from_MatMul(M*N*P)
    assert recognize_matrix_expression(cg) == M*N*P

    cg = parse_indexed_expression((M*N.T*P)[i,j])
    assert recognize_matrix_expression(cg) == M*N.T*P
    cg = CodegenArrayContraction.from_MatMul(M*N.T*P)
    assert recognize_matrix_expression(cg) == M*N.T*P

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(M,N,P,Q), (1, 2), (5, 6))
    assert recognize_matrix_expression(cg) == [M*N, P*Q]

    expr = -2*M*N
    elem = expr[i, j]
    cg = parse_indexed_expression(elem)
    assert recognize_matrix_expression(cg) == -2*M*N
示例#4
0
def test_codegen_array_contraction_construction():
    cg = CodegenArrayContraction(A)
    assert cg == A

    s = Sum(A[i]*B[i], (i, 0, 3))
    cg = parse_indexed_expression(s)
    assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (0, 1))

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (1, 0))
    assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(A, B), (0, 1))

    expr = M*N
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))
    assert CodegenArrayContraction.from_MatMul(expr) == result
    elem = expr[i, j]
    assert parse_indexed_expression(elem) == result

    expr = M*N*M
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, M), (1, 2), (3, 4))
    assert CodegenArrayContraction.from_MatMul(expr) == result
    elem = expr[i, j]
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N), (1, 4), (2, 5))
    cg = parse_indexed_expression(elem)
    cg = cg.sort_args_by_name()
    assert cg == result
示例#5
0
def test_codegen_einsum():
    if not np:
        skip("NumPy not installed")

    M = MatrixSymbol("M", 2, 2)
    N = MatrixSymbol("N", 2, 2)

    cg = CodegenArrayContraction.from_MatMul(M*N)
    f = lambdify((M, N), cg, 'numpy')

    ma = np.matrix([[1, 2], [3, 4]])
    mb = np.matrix([[1,-2], [-1, 3]])
    assert (f(ma, mb) == ma*mb).all()
示例#6
0
def test_codegen_einsum():
    if not np:
        skip("NumPy not installed")

    M = MatrixSymbol("M", 2, 2)
    N = MatrixSymbol("N", 2, 2)

    cg = CodegenArrayContraction.from_MatMul(M*N)
    f = lambdify((M, N), cg, 'numpy')

    ma = np.matrix([[1, 2], [3, 4]])
    mb = np.matrix([[1,-2], [-1, 3]])
    assert (f(ma, mb) == ma*mb).all()
示例#7
0
def test_codegen_einsum():
    if not tf:
        skip("TensorFlow not installed")

    session = tf.Session()

    M = MatrixSymbol("M", 2, 2)
    N = MatrixSymbol("N", 2, 2)

    cg = CodegenArrayContraction.from_MatMul(M * N)
    f = lambdify((M, N), cg, 'tensorflow')

    ma = tf.constant([[1, 2], [3, 4]])
    mb = tf.constant([[1, -2], [-1, 3]])
    y = session.run(f(ma, mb))
    c = session.run(tf.matmul(ma, mb))
    assert (y == c).all()
示例#8
0
def test_codegen_einsum():
    if not tf:
        skip("TensorFlow not installed")

    session = tf.Session()

    M = MatrixSymbol("M", 2, 2)
    N = MatrixSymbol("N", 2, 2)

    cg = CodegenArrayContraction.from_MatMul(M*N)
    f = lambdify((M, N), cg, 'tensorflow')

    ma = tf.constant([[1, 2], [3, 4]])
    mb = tf.constant([[1,-2], [-1, 3]])
    y = session.run(f(ma, mb))
    c = session.run(tf.matmul(ma, mb))
    assert (y == c).all()
示例#9
0
def test_codegen_einsum():
    if not tf:
        skip("TensorFlow not installed")

    graph = tf.Graph()
    with graph.as_default():
        session = tf.compat.v1.Session(graph=graph)

        M = MatrixSymbol("M", 2, 2)
        N = MatrixSymbol("N", 2, 2)

        cg = CodegenArrayContraction.from_MatMul(M * N)
        f = lambdify((M, N), cg, "tensorflow")

        ma = tf.constant([[1, 2], [3, 4]])
        mb = tf.constant([[1, -2], [-1, 3]])
        y = session.run(f(ma, mb))
        c = session.run(tf.matmul(ma, mb))
        assert (y == c).all()