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
