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