def test_codegen_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
assert recognize_matrix_expression(expr) == M + Transpose(M)
expr = M[i,j] + N[i,j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N
expr = M[i,j] + N[j,i]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N.T
expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M))
expr = (M*N*P)[i, j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*N*P
expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*N*P
expr = Sum((P[j, m] + P[m, j])*(M[i,j]*N[m,n] + N[i,j]*M[m,n]), (j, 0, k-1), (m, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*P*N + M*P.T*N + N*P*M + N*P.T*M
def test_codegen_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
rec = _recognize_matrix_expression(expr)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [M])])
assert _unfold_recognized_expr(rec) == M + Transpose(M)
expr = M[i, j] + N[i, j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, N])
assert _unfold_recognized_expr(rec) == M + N
expr = M[i, j] + N[j, i]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [N])])
assert _unfold_recognized_expr(rec) == M + N.T
expr = M[i, j] * N[k, l] + N[i, j] * M[k, l]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(
MatAdd, [_RecognizeMatMulLines([M, N]),
_RecognizeMatMulLines([N, M])])
#assert _unfold_recognized_expr(rec) == TensorProduct(M, N) + TensorProduct(N, M) maybe?
expr = (M * N * P)[i, j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatMulLines([_RecognizeMatOp(MatMul, [M, N, P])])
assert _unfold_recognized_expr(rec) == M * N * P
expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatMul, [M, N, P])
assert _unfold_recognized_expr(rec) == M * N * P
expr = Sum((P[j, m] + P[m, j]) * (M[i, j] * N[m, n] + N[i, j] * M[m, n]),
(j, 0, k - 1), (m, 0, k - 1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [
_RecognizeMatOp(MatMul, [
M,
_RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), N
]),
_RecognizeMatOp(MatMul, [
N,
_RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), M
])
])
assert _unfold_recognized_expr(
rec) == M * (P + P.T) * N + N * (P + P.T) * M
def test_codegen_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
rec = _recognize_matrix_expression(expr)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [M])])
assert _unfold_recognized_expr(rec) == M + Transpose(M)
expr = M[i,j] + N[i,j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, N])
assert _unfold_recognized_expr(rec) == M + N
expr = M[i,j] + N[j,i]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [N])])
assert _unfold_recognized_expr(rec) == M + N.T
expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [_RecognizeMatMulLines([M, N]), _RecognizeMatMulLines([N, M])])
#assert _unfold_recognized_expr(rec) == TensorProduct(M, N) + TensorProduct(N, M) maybe?
expr = (M*N*P)[i, j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatMulLines([_RecognizeMatOp(MatMul, [M, N, P])])
assert _unfold_recognized_expr(rec) == M*N*P
expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatMul, [M, N, P])
assert _unfold_recognized_expr(rec) == M*N*P
expr = Sum((P[j, m] + P[m, j])*(M[i,j]*N[m,n] + N[i,j]*M[m,n]), (j, 0, k-1), (m, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [
_RecognizeMatOp(MatMul, [M, _RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), N]),
_RecognizeMatOp(MatMul, [N, _RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), M])
])
assert _unfold_recognized_expr(rec) == M*(P + P.T)*N + N*(P + P.T)*M
def test_codegen_array_parse():
expr = M[i, j]
assert _codegen_array_parse(expr) == (M, (i, j))
expr = M[i, j] * N[k, l]
assert _codegen_array_parse(expr) == (CodegenArrayTensorProduct(M, N),
(i, j, k, l))
expr = M[i, j] * N[j, k]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1))
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, N),
(i, j))
expr = M[i, j] + N[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(N, Permutation([1, 0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(M, Permutation([1, 0]))), (i, j))
expr = (M * N * P)[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _codegen_array_parse(expr)
assert ret1 == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P),
(1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j) * M[i, k]
assert _codegen_array_parse(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * M[i, l]
assert _codegen_array_parse(expr) == (M, ({i, j, k}, l))
expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, k})))
expr = KroneckerDelta(j, m) * KroneckerDelta(
m, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, m, k})))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * KroneckerDelta(
k, m) * M[i, 0] * KroneckerDelta(m, n)
assert _codegen_array_parse(expr) == (M, ({i, j, k, m, n}, 0))
expr = M[i, i]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(M,
(0, 1)), (i, ))
def test_codegen_array_parse():
expr = M[i, j]
assert _codegen_array_parse(expr) == (M, (i, j))
expr = M[i, j]*N[k, l]
assert _codegen_array_parse(expr) == (CodegenArrayTensorProduct(M, N), (i, j, k, l))
expr = M[i, j]*N[j, k]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j]*N[j, k], (j, 0, k-1))
assert _codegen_array_parse(expr) == (CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, N), (i, j))
expr = M[i, j] + N[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(N, Permutation([1,0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, Permutation([1,0]))), (i, j))
expr = (M*N*P)[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _codegen_array_parse(expr)
assert ret1 == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j)*M[i, k]
assert _codegen_array_parse(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(i, j)*KroneckerDelta(j, k)*M[i, l]
assert _codegen_array_parse(expr) == (M, ({i, j, k}, l))
expr = KroneckerDelta(j, k)*(M[i, j]*N[k, l] + N[i, j]*M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), Permutation(0, 2)(1, 3))
), (1, 2)), (i, l, frozenset({j, k})))
expr = KroneckerDelta(j, m)*KroneckerDelta(m, k)*(M[i, j]*N[k, l] + N[i, j]*M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), Permutation(0, 2)(1, 3))
), (1, 2)), (i, l, frozenset({j, m, k})))
expr = KroneckerDelta(i, j)*KroneckerDelta(j, k)*KroneckerDelta(k,m)*M[i, 0]*KroneckerDelta(m, n)
assert _codegen_array_parse(expr) == (M, ({i,j,k,m,n}, 0))
expr = M[i, i]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(M, (0, 1)), (i,))