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