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 = parse_matrix_expression(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 = parse_matrix_expression(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 parse_matrix_expression(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 parse_matrix_expression(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_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 = parse_matrix_expression(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 = parse_matrix_expression(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) == CodegenArrayTensorProduct( 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 a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) c = MatrixSymbol("c", k, 1) cg = CodegenArrayPermuteDims( CodegenArrayContraction( CodegenArrayTensorProduct( a, CodegenArrayElementwiseAdd( CodegenArrayTensorProduct(b, c), CodegenArrayTensorProduct(c, b), )), (2, 4)), [0, 1, 3, 2]) assert recognize_matrix_expression(cg) == a * (b.T * c + c.T * b) za = ZeroArray(m, n) assert recognize_matrix_expression(za) == ZeroMatrix(m, n) cg = CodegenArrayTensorProduct(3, M) assert recognize_matrix_expression(cg) == 3 * M
def test_special_matrices(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) expr = a.T*b elem = expr[0, 0] cg = parse_indexed_expression(elem) assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(a, b), (0, 2)) assert recognize_matrix_expression(cg) == a.T*b
def test_special_matrices(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) expr = a.T*b elem = expr[0, 0] cg = parse_indexed_expression(elem) assert cg == CodegenArrayContraction(CodegenArrayTensorProduct(a, b), (0, 2)) assert recognize_matrix_expression(cg) == a.T*b
def test_codegen_array_parse_out_of_bounds(): expr = Sum(M[i, i], (i, 0, 4)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, i], (i, 0, k)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, i], (i, 1, k - 1)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j] * N[j, m], (j, 0, 4)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j] * N[j, m], (j, 0, k)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j] * N[j, m], (j, 1, k - 1)) raises(ValueError, lambda: parse_indexed_expression(expr))
def test_codegen_array_parse_out_of_bounds(): expr = Sum(M[i, i], (i, 0, 4)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, i], (i, 0, k)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, i], (i, 1, k-1)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j]*N[j,m], (j, 0, 4)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j]*N[j,m], (j, 0, k)) raises(ValueError, lambda: parse_indexed_expression(expr)) expr = Sum(M[i, j]*N[j,m], (j, 1, k-1)) raises(ValueError, lambda: parse_indexed_expression(expr))