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
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
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
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
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()
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()
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()
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()
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()