def test_arrayexpr_derivatives1(): res = array_derive(X, X) assert res == CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, I), [0, 2, 1, 3]) cg = CodegenArrayTensorProduct(A, X, B) res = array_derive(cg, X) assert res == CodegenArrayPermuteDims( CodegenArrayTensorProduct(I, A, I, B), [0, 4, 2, 3, 1, 5, 6, 7]) cg = CodegenArrayContraction(X, (0, 1)) res = array_derive(cg, X) assert res == CodegenArrayContraction(CodegenArrayTensorProduct(I, I), (1, 3)) cg = CodegenArrayDiagonal(X, (0, 1)) res = array_derive(cg, X) assert res == CodegenArrayDiagonal(CodegenArrayTensorProduct(I, I), (1, 3)) cg = ElementwiseApplyFunction(sin, X) res = array_derive(cg, X) assert res.dummy_eq( CodegenArrayDiagonal( CodegenArrayTensorProduct(ElementwiseApplyFunction(cos, X), I, I), (0, 3), (1, 5)))
def test_array_expr_zero_array(): za1 = ZeroArray(k, l, m, n) zm1 = ZeroMatrix(m, n) za2 = ZeroArray(k, m, m, n) zm2 = ZeroMatrix(m, m) zm3 = ZeroMatrix(k, k) assert CodegenArrayTensorProduct(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n) assert CodegenArrayTensorProduct(M, N, zm1) == ZeroArray(k, k, k, k, m, n) assert CodegenArrayContraction(za1, (3,)) == ZeroArray(k, l, m) assert CodegenArrayContraction(zm1, (1,)) == ZeroArray(m) assert CodegenArrayContraction(za2, (1, 2)) == ZeroArray(k, n) assert CodegenArrayContraction(zm2, (0, 1)) == 0 assert CodegenArrayDiagonal(za2, (1, 2)) == ZeroArray(k, n, m) assert CodegenArrayDiagonal(zm2, (0, 1)) == ZeroArray(m) assert CodegenArrayPermuteDims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k) assert CodegenArrayPermuteDims(zm1, [1, 0]) == ZeroArray(n, m) assert CodegenArrayElementwiseAdd(za1) == za1 assert CodegenArrayElementwiseAdd(zm1) == ZeroArray(m, n) tp1 = CodegenArrayTensorProduct(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n)) assert CodegenArrayElementwiseAdd(tp1, za1) == tp1 tp2 = CodegenArrayTensorProduct(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n)) assert CodegenArrayElementwiseAdd(tp1, za1, tp2) == CodegenArrayElementwiseAdd(tp1, tp2) assert CodegenArrayElementwiseAdd(M, zm3) == M assert CodegenArrayElementwiseAdd(M, N, zm3) == CodegenArrayElementwiseAdd(M, N)
def test_contraction_permutation_mix(): Me = M.subs(k, 3).as_explicit() Ne = N.subs(k, 3).as_explicit() cg1 = CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), Permutation([0, 2, 1, 3])), (2, 3)) cg2 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 3)) assert cg1 == cg2 assert recognize_matrix_expression(cg2) == M * N.T cge1 = tensorcontraction( permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3)) cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3)) assert cge1 == cge2 cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), Permutation([0, 1, 3, 2])) cg2 = CodegenArrayTensorProduct( M, CodegenArrayPermuteDims(N, Permutation([1, 0]))) assert cg1 == cg2 assert recognize_matrix_expression(cg1) == CodegenArrayTensorProduct( M, N.T) assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct( M, N.T) cg1 = CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5)) cg2 = CodegenArrayContraction( CodegenArrayTensorProduct( M, N, P, CodegenArrayPermuteDims(Q, Permutation([1, 0]))), (1, 5), (2, 3)) assert cg1 == cg2 assert recognize_matrix_expression(cg1) == CodegenArrayTensorProduct( M * P.T * Trace(N), Q.T) assert recognize_matrix_expression(cg2) == CodegenArrayTensorProduct( M * P.T * Trace(N), Q.T) cg1 = CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])), (0, 1), (2, 6), (3, 7)) cg2 = CodegenArrayPermuteDims( CodegenArrayContraction(CodegenArrayTensorProduct(M, P, Q, N), (0, 1), (2, 3), (4, 7)), [1, 0]) assert cg1 == cg2 cg1 = CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])), (0, 1), (2, 6), (3, 7)) cg2 = CodegenArrayPermuteDims( CodegenArrayContraction( CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), N, P, Q), (0, 1), (3, 6), (4, 5)), Permutation([1, 0])) assert cg1 == cg2
def test_remove_trivial_dims(): # Tensor Product: assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b)) == (a * b.T, [1, 3]) assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b)) == (a * b.T, [0, 3]) assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b.T)) == (a * b.T, [1, 2]) assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b.T)) == (a * b.T, [0, 2]) assert _remove_trivial_dims(CodegenArrayTensorProduct( I, a.T, b.T)) == (a * b.T, [0, 1, 2, 4]) assert _remove_trivial_dims(CodegenArrayTensorProduct( a.T, I, b.T)) == (a * b.T, [0, 2, 3, 4]) assert _remove_trivial_dims(CodegenArrayTensorProduct(a, I)) == (a, [2, 3]) assert _remove_trivial_dims(CodegenArrayTensorProduct(I, a)) == (a, [0, 1]) assert _remove_trivial_dims(CodegenArrayTensorProduct( a.T, b.T, c, d)) == (CodegenArrayTensorProduct(a * b.T, c * d.T), [0, 2, 5, 7]) assert _remove_trivial_dims(CodegenArrayTensorProduct( a.T, I, b.T, c, d, I)) == (CodegenArrayTensorProduct(a * b.T, c * d.T, I), [0, 2, 3, 4, 7, 9]) # Addition: cg = CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(a, b), CodegenArrayTensorProduct(c, d)) assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3]) # Permute Dims: cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, b), Permutation(3)(1, 2)) assert _remove_trivial_dims(cg) == (a * b.T, [2, 3]) cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, I, b), Permutation(5)(1, 2, 3, 4)) assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5]) cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, b, a), Permutation(5)(1, 2, 4, 5, 3)) assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5]) # Diagonal: cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, []) # Contraction: cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, [])
def test_arrayexpr_derivatives1(): res = array_derive(X, X) assert res == CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, I), [0, 2, 1, 3]) cg = CodegenArrayTensorProduct(A, X, B) res = array_derive(cg, X) assert res == CodegenArrayPermuteDims( CodegenArrayTensorProduct(I, A, I, B), [0, 4, 2, 3, 1, 5, 6, 7])
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_normalize_diagonal_permutedims(): tp = CodegenArrayTensorProduct(M, Q, N, P) expr = CodegenArrayDiagonal( CodegenArrayPermuteDims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7), (0, 3)) result = CodegenArrayDiagonal(tp, (2, 6, 7), (3, 5), (0, 4)) assert expr == result tp = CodegenArrayTensorProduct(M, N, P, Q) expr = CodegenArrayDiagonal(CodegenArrayPermuteDims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4)) result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, P, N, Q), (3, 4, 5), (1, 2)) assert expr == result
def test_parsing_of_matrix_expressions(): expr = M * N assert _parse_matrix_expression(expr) == CodegenArrayContraction( CodegenArrayTensorProduct(M, N), (1, 2)) expr = Transpose(M) assert _parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0]) expr = M * Transpose(N) assert _parse_matrix_expression(expr) == CodegenArrayContraction( CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2))
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 _(expr: CodegenArrayTensorProduct, x: Expr): args = expr.args addend_list = [] for i, arg in enumerate(expr.args): darg = array_derive(arg, x) if darg == 0: continue args_prev = args[:i] args_succ = args[i + 1:] shape_prev = reduce(operator.add, map(get_shape, args_prev), ()) shape_succ = reduce(operator.add, map(get_shape, args_succ), ()) addend = CodegenArrayTensorProduct(*args_prev, darg, *args_succ) tot1 = len(get_shape(x)) tot2 = tot1 + len(shape_prev) tot3 = tot2 + len(get_shape(arg)) tot4 = tot3 + len(shape_succ) perm = [i for i in range(tot1, tot2)] + \ [i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \ [i for i in range(tot3, tot4)] addend = CodegenArrayPermuteDims(addend, _af_invert(perm)) addend_list.append(addend) if len(addend_list) == 1: return addend_list[0] elif len(addend_list) == 0: return S.Zero else: return CodegenArrayElementwiseAdd(*addend_list)
def _(expr: Inverse, x: Expr): mat = expr.I dexpr = array_derive(mat, x) tp = CodegenArrayTensorProduct(-expr, dexpr, expr) mp = CodegenArrayContraction(tp, (1, 4), (5, 6)) pp = CodegenArrayPermuteDims(mp, [1, 2, 0, 3]) return pp
def test_nested_permutations(): cg = CodegenArrayPermuteDims(CodegenArrayPermuteDims(M, (1, 0)), (1, 0)) assert cg == M times = 3 plist1 = [list(range(6)) for i in range(times)] plist2 = [list(range(6)) for i in range(times)] for i in range(times): random.shuffle(plist1[i]) random.shuffle(plist2[i]) plist1.append([2, 5, 4, 1, 0, 3]) plist2.append([3, 5, 0, 4, 1, 2]) plist1.append([2, 5, 4, 0, 3, 1]) plist2.append([3, 0, 5, 1, 2, 4]) plist1.append([5, 4, 2, 0, 3, 1]) plist2.append([4, 5, 0, 2, 3, 1]) Me = M.subs(k, 3).as_explicit() Ne = N.subs(k, 3).as_explicit() Pe = P.subs(k, 3).as_explicit() cge = tensorproduct(Me, Ne, Pe) for permutation_array1, permutation_array2 in zip(plist1, plist2): p1 = Permutation(permutation_array1) p2 = Permutation(permutation_array2) cg = CodegenArrayPermuteDims( CodegenArrayPermuteDims( CodegenArrayTensorProduct(M, N, P), p1), p2 ) result = CodegenArrayPermuteDims( CodegenArrayTensorProduct(M, N, P), p2*p1 ) assert cg == result # Check that `permutedims` behaves the same way with explicit-component arrays: result1 = permutedims(permutedims(cge, p1), p2) result2 = permutedims(cge, p2*p1) assert result1 == result2
def test_codegen_extra(): if not np: skip("NumPy not installed") M = MatrixSymbol("M", 2, 2) N = MatrixSymbol("N", 2, 2) P = MatrixSymbol("P", 2, 2) Q = MatrixSymbol("Q", 2, 2) ma = np.matrix([[1, 2], [3, 4]]) mb = np.matrix([[1, -2], [-1, 3]]) mc = np.matrix([[2, 0], [1, 2]]) md = np.matrix([[1, -1], [4, 7]]) cg = CodegenArrayTensorProduct(M, N) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all() cg = CodegenArrayElementwiseAdd(M, N) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == ma + mb).all() cg = CodegenArrayElementwiseAdd(M, N, P) f = lambdify((M, N, P), cg, 'numpy') assert (f(ma, mb, mc) == ma + mb + mc).all() cg = CodegenArrayElementwiseAdd(M, N, P, Q) f = lambdify((M, N, P, Q), cg, 'numpy') assert (f(ma, mb, mc, md) == ma + mb + mc + md).all() cg = CodegenArrayPermuteDims(M, [1, 0]) f = lambdify((M, ), cg, 'numpy') assert (f(ma) == ma.T).all() cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 2, 3, 0]) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.transpose(np.einsum(ma, [0, 1], mb, [2, 3]), (1, 2, 3, 0))).all() cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2)) f = lambdify((M, N), cg, 'numpy') assert (f(ma, mb) == np.diagonal(np.einsum(ma, [0, 1], mb, [2, 3]), axis1=1, axis2=2)).all()
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_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_parsing_of_matrix_expressions(): expr = M * N assert parse_matrix_expression(expr) == CodegenArrayContraction( CodegenArrayTensorProduct(M, N), (1, 2)) expr = Transpose(M) assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0]) expr = M * Transpose(N) assert parse_matrix_expression(expr) == CodegenArrayContraction( CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2)) expr = 3 * M * N res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr == rexpr expr = 3 * M + N * M.T * M + 4 * k * N res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr == rexpr expr = Inverse(M) * N rexpr = recognize_matrix_expression(parse_matrix_expression(expr)) assert expr == rexpr expr = M**2 rexpr = recognize_matrix_expression(parse_matrix_expression(expr)) assert expr == rexpr expr = M * (2 * N + 3 * M) res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr.expand() == rexpr.doit() expr = Trace(M) result = CodegenArrayContraction(M, (0, 1)) assert parse_matrix_expression(expr) == result
def test_codegen_array_doit(): M = MatrixSymbol("M", 2, 2) N = MatrixSymbol("N", 2, 2) P = MatrixSymbol("P", 2, 2) Q = MatrixSymbol("Q", 2, 2) M = M.as_explicit() N = N.as_explicit() P = P.as_explicit() Q = Q.as_explicit() expr = CodegenArrayTensorProduct(M, N, P, Q) assert expr.doit() == tensorproduct(M, N, P, Q) expr2 = CodegenArrayContraction(expr, (0, 1)) assert expr2.doit() == tensorcontraction(tensorproduct(M, N, P, Q), (0, 1)) expr2 = CodegenArrayDiagonal(expr, (0, 1)) #assert expr2 = ... # TODO: not implemented expr = CodegenArrayTensorProduct(M, N) exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0]) assert exprp.doit() == permutedims(tensorproduct(M, N), [2, 1, 3, 0]) expr = CodegenArrayElementwiseAdd(M, N) assert expr.doit() == M + N
def test_recognize_expression_implicit_mul(): cg = CodegenArrayTensorProduct(a, b) assert recognize_matrix_expression(cg) == a*b.T cg = CodegenArrayTensorProduct(a, I, b) assert recognize_matrix_expression(cg) == a*b.T cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I), (1, 2)) assert recognize_matrix_expression(cg) == I cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, Identity(1)), [0, 2, 1, 3]) assert recognize_matrix_expression(cg) == I
def test_codegen_array_shape(): expr = CodegenArrayTensorProduct(M, N, P, Q) assert expr.shape == (k, k, k, k, k, k, k, k) Z = MatrixSymbol("Z", m, n) expr = CodegenArrayTensorProduct(M, Z) assert expr.shape == (k, k, m, n) expr2 = CodegenArrayContraction(expr, (0, 1)) assert expr2.shape == (m, n) expr2 = CodegenArrayDiagonal(expr, (0, 1)) assert expr2.shape == (m, n, k) exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0]) assert exprp.shape == (m, k, n, k) expr3 = CodegenArrayTensorProduct(N, Z) expr2 = CodegenArrayElementwiseAdd(expr, expr3) assert expr2.shape == (k, k, m, n) # Contraction along axes with discordant dimensions: raises(ValueError, lambda: CodegenArrayContraction(expr, (1, 2))) # Also diagonal needs the same dimensions: raises(ValueError, lambda: CodegenArrayDiagonal(expr, (1, 2)))
def test_permute_tensor_product(): cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 1, 0, 5, 4, 6, 7])) cg2 = CodegenArrayTensorProduct(N, CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(P, [1, 0]), Q) assert cg1 == cg2 # TODO: reverse operation starting with `CodegenArrayPermuteDims` and getting down to `bb`... cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 4, 5, 0, 1, 6, 7])) cg2 = CodegenArrayTensorProduct(N, P, M, Q) assert cg1 == cg2 cg1 = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P, Q), Permutation([2, 3, 4, 6, 5, 7, 0, 1])) assert cg1.expr == CodegenArrayTensorProduct(N, P, Q, M) assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7]) cg1 = CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]) cg2 = CodegenArrayPermuteDims( CodegenArrayContraction(CodegenArrayTensorProduct(Q, Q, N, M), (3, 5, 6)), [0, 2, 3, 1, 4]) assert cg1 == cg2 cg1 = CodegenArrayContraction( CodegenArrayContraction( CodegenArrayContraction( CodegenArrayContraction( CodegenArrayPermuteDims( CodegenArrayTensorProduct(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]), [1, 3, 4]), [1]), [0]) cg2 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, Q, Q), (0, 3, 5), (1, 4, 7), (2, ), (6, )) assert cg1 == cg2
def _(expr: MatrixSymbol, x: Expr): m, n = expr.shape if expr == x: return CodegenArrayPermuteDims( CodegenArrayTensorProduct(Identity(m), Identity(n)), [0, 2, 1, 3]) return ZeroArray(*(x.shape + expr.shape))
def _(expr: CodegenArrayPermuteDims, x: Expr): de = array_derive(expr.expr, x) perm = [0, 1] + [i + 2 for i in expr.permutation.array_form] return CodegenArrayPermuteDims(de, perm)
def _(expr: Transpose, x: Expr): # D(A.T, A) ==> (m,n,i,j) ==> D(A_ji, A_mn) = d_mj d_ni # D(B.T, A) ==> (m,n,i,j) ==> D(B_ji, A_mn) fd = array_derive(expr.arg, x) return CodegenArrayPermuteDims(fd, [0, 1, 3, 2])
def test_codegen_permutedims_sink(): cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [0, 1, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct( M, CodegenArrayPermuteDims(N, [1, 0])) assert recognize_matrix_expression(sunk) == [M, N.T] cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct( CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0])) assert recognize_matrix_expression(sunk) == [M.T, N.T] cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [3, 2, 1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct( CodegenArrayPermuteDims(N, [1, 0]), CodegenArrayPermuteDims(M, [1, 0])) assert recognize_matrix_expression(sunk) == [N.T, M.T] cg = CodegenArrayPermuteDims( CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)), [1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [[0, 3]]), (1, 2)) cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct( CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0])) cg = CodegenArrayPermuteDims( CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayContraction( CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4))
def permutedims(expr, perm): """ Permutes the indices of an array. Parameter specifies the permutation of the indices. Examples ======== >>> from sympy.abc import x, y, z, t >>> from sympy import sin >>> from sympy import Array, permutedims >>> a = Array([[x, y, z], [t, sin(x), 0]]) >>> a [[x, y, z], [t, sin(x), 0]] >>> permutedims(a, (1, 0)) [[x, t], [y, sin(x)], [z, 0]] If the array is of second order, ``transpose`` can be used: >>> from sympy import transpose >>> transpose(a) [[x, t], [y, sin(x)], [z, 0]] Examples on higher dimensions: >>> b = Array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) >>> permutedims(b, (2, 1, 0)) [[[1, 5], [3, 7]], [[2, 6], [4, 8]]] >>> permutedims(b, (1, 2, 0)) [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] ``Permutation`` objects are also allowed: >>> from sympy.combinatorics import Permutation >>> permutedims(b, Permutation([1, 2, 0])) [[[1, 5], [2, 6]], [[3, 7], [4, 8]]] """ from sympy.tensor.array import SparseNDimArray from sympy.tensor.array.expressions.array_expressions import _ArrayExpr from sympy.codegen.array_utils import _CodegenArrayAbstract, CodegenArrayPermuteDims if isinstance(expr, (_ArrayExpr, _CodegenArrayAbstract)): return CodegenArrayPermuteDims(expr, perm) if not isinstance(expr, NDimArray): expr = ImmutableDenseNDimArray(expr) from sympy.combinatorics import Permutation if not isinstance(perm, Permutation): perm = Permutation(list(perm)) if perm.size != expr.rank(): raise ValueError("wrong permutation size") # Get the inverse permutation: iperm = ~perm new_shape = perm(expr.shape) if isinstance(expr, SparseNDimArray): return type(expr)({tuple(perm(expr._get_tuple_index(k))): v for k, v in expr._sparse_array.items()}, new_shape) indices_span = perm([range(i) for i in expr.shape]) new_array = [None]*len(expr) for i, idx in enumerate(itertools.product(*indices_span)): t = iperm(idx) new_array[i] = expr[t] return type(expr)(new_array, new_shape)
def test_codegen_permutedims_sink(): cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [0, 1, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])) assert recognize_matrix_expression(sunk) == [M, N.T] cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0])) assert recognize_matrix_expression(sunk) == [M.T, N.T] cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [3, 2, 1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct(CodegenArrayPermuteDims(N, [1, 0]), CodegenArrayPermuteDims(M, [1, 0])) assert recognize_matrix_expression(sunk) == [N.T, M.T] cg = CodegenArrayPermuteDims(CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)), [1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [[0, 3]]), (1, 2)) cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 0, 3, 2]) sunk = cg.nest_permutation() assert sunk == CodegenArrayTensorProduct(CodegenArrayPermuteDims(M, [1, 0]), CodegenArrayPermuteDims(N, [1, 0])) cg = CodegenArrayPermuteDims(CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), [1, 0]) sunk = cg.nest_permutation() assert sunk == CodegenArrayContraction(CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N, P), [[0, 5]]), (1, 2), (3, 4)) sunk2 = sunk.expr.nest_permutation()
def test_recognize_products_support_function(): A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) D = MatrixSymbol("D", k, k) X = MatrixSymbol("X", k, k) Y = MatrixSymbol("Y", k, k) assert _support_function_tp1_recognize([], [2 * k]) == 2 * k assert _support_function_tp1_recognize([(1, 2)], [A, 2 * k, B, 3]) == 6 * k * A * B assert _support_function_tp1_recognize([(0, 3), (1, 2)], [A, B]) == Trace(A * B) assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T assert _support_function_tp1_recognize( [(1, 2), (5, 6)], [A, B, C, D]) == CodegenArrayTensorProduct(A * B, C * D) assert _support_function_tp1_recognize( [(1, 4), (3, 6)], [A, B, C, D]) == CodegenArrayPermuteDims( CodegenArrayTensorProduct(A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize([(0, 3), (1, 4)], [A, B, C]) == B * A * C assert _support_function_tp1_recognize( [(9, 10), (1, 2), (5, 6), (3, 4), (7, 8)], [X, Y, A, B, C, D]) == X * Y * A * B * C * D assert _support_function_tp1_recognize( [(9, 10), (1, 2), (5, 6), (3, 4)], [X, Y, A, B, C, D]) == CodegenArrayTensorProduct(X * Y * A * B, C * D) assert _support_function_tp1_recognize( [(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == CodegenArrayPermuteDims( CodegenArrayTensorProduct(X * B.T, Y * C, D * A), [0, 2, 5, 1, 3, 4]) assert _support_function_tp1_recognize( [(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == CodegenArrayPermuteDims( CodegenArrayTensorProduct(Trace(X) * A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [A, A, B, C, D]) == A**2 * B * C * D assert _support_function_tp1_recognize( [(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D assert _support_function_tp1_recognize( [(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == CodegenArrayPermuteDims( CodegenArrayTensorProduct(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5]) assert _support_function_tp1_recognize( [(1, 4), (3, 6)], [A, B, C, D]) == CodegenArrayPermuteDims( CodegenArrayTensorProduct(A * C, B * D), [0, 2, 1, 3]) assert _support_function_tp1_recognize( [(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C * X.T * B * A * D assert _support_function_tp1_recognize( [(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C * X.T * B * A * D
def test_recognize_diagonalized_vectors(): a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) A = MatrixSymbol("A", k, k) B = MatrixSymbol("B", k, k) C = MatrixSymbol("C", k, k) X = MatrixSymbol("X", k, k) x = MatrixSymbol("x", k, 1) I1 = Identity(1) I = Identity(k) # Check matrix recognition over trivial dimensions: cg = CodegenArrayTensorProduct(a, b) assert recognize_matrix_expression(cg) == a * b.T cg = CodegenArrayTensorProduct(I1, a, b) assert recognize_matrix_expression(cg) == a * b.T # Recognize trace inside a tensor product: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B, C), (0, 3), (1, 2)) assert recognize_matrix_expression(cg) == Trace(A * B) * C # Transform diagonal operator to contraction: cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (1, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (1, 3)) assert recognize_matrix_expression(cg) == A * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayPermuteDims( CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]) assert recognize_matrix_expression(cg) == b.T * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(A, OneArray(1), DiagMatrix(a)), (0, 3)) assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2), (3, 5)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(I, OneArray(1), I1, DiagMatrix(x)), (0, 5)) assert recognize_matrix_expression(cg) == DiagMatrix(x) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2), (5, 6)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayDiagonal( CodegenArrayContraction( CodegenArrayTensorProduct(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6)) # TODO: not yet working # assert recognize_matrix_expression(cg) cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I1), (1, 2)) assert isinstance(cg, CodegenArrayDiagonal) assert cg.diagonal_indices == ((1, 2), ) assert recognize_matrix_expression(cg) == x cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I), (0, 2)) assert _array_diag2contr_diagmatrix(cg) == CodegenArrayContraction( CodegenArrayTensorProduct(OneArray(1), I, DiagMatrix(x)), (1, 3)) assert recognize_matrix_expression(cg).doit() == DiagMatrix(x) raises(ValueError, lambda: CodegenArrayDiagonal(x, (1, ))) # Ignore identity matrices with contractions: cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, A, I, I), (0, 2), (1, 3), (5, 7)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg) == Trace(A) * I cg = CodegenArrayContraction(CodegenArrayTensorProduct(Trace(A) * I, I, I), (1, 5), (3, 4)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg).doit() == Trace(A) * I # Add DiagMatrix when required: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2)) assert cg.split_multiple_contractions() == cg assert recognize_matrix_expression(cg) == A * a cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), B), (1, 2), (3, 4)) assert recognize_matrix_expression(cg) == A * DiagMatrix(a) * B cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), B), (0, 2), (3, 4)) assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) * B cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B), (0, 2, 4, 7, 9)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(A, DiagMatrix(a), DiagMatrix(b), DiagMatrix(a), B), (0, 2), (3, 4), (5, 7), (6, 9)) assert recognize_matrix_expression( cg).doit() == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T cg = CodegenArrayContraction(CodegenArrayTensorProduct(I1, I1, I1), (1, 2, 4)) assert cg.split_multiple_contractions() == CodegenArrayContraction( CodegenArrayTensorProduct(I1, I1, I1), (1, 2), (3, 4)) assert recognize_matrix_expression(cg).doit() == Identity(1) cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I, I, I, A), (1, 2, 8), (5, 6, 9)) assert recognize_matrix_expression( cg.split_multiple_contractions()).doit() == A cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8)) expected = CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), DiagMatrix(a), C, A, B), (0, 4), (1, 7), (2, 5), (3, 8)) assert cg.split_multiple_contractions() == expected assert recognize_matrix_expression( cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B cg = CodegenArrayContraction( CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) assert cg.split_multiple_contractions().dummy_eq( CodegenArrayContraction( CodegenArrayTensorProduct((a.T * b).applyfunc(cos), I1, I1, a, b), (0, 2), (1, 4), (3, 7), (5, 9))) assert recognize_matrix_expression(cg).doit().dummy_eq( MatMul(a, (a.T * b).applyfunc(cos), b.T)) cg = CodegenArrayContraction( CodegenArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) assert cg.split_multiple_contractions().dummy_eq( CodegenArrayContraction( CodegenArrayTensorProduct(DiagMatrix(a), (A * X * b).applyfunc(cos), A.T, b, b.T), (0, 2), (1, 5), (3, 7, 8))) # assert recognize_matrix_expression(cg) # Check no overlap of lines: cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = CodegenArrayContraction(CodegenArrayTensorProduct(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
def test_codegen_extra(): if not tf: skip("TensorFlow not installed") graph = tf.Graph() with graph.as_default(): session = tf.compat.v1.Session() M = MatrixSymbol("M", 2, 2) N = MatrixSymbol("N", 2, 2) P = MatrixSymbol("P", 2, 2) Q = MatrixSymbol("Q", 2, 2) ma = tf.constant([[1, 2], [3, 4]]) mb = tf.constant([[1, -2], [-1, 3]]) mc = tf.constant([[2, 0], [1, 2]]) md = tf.constant([[1, -1], [4, 7]]) cg = CodegenArrayTensorProduct(M, N) assert tensorflow_code(cg) == \ 'tensorflow.linalg.einsum("ab,cd", M, N)' f = lambdify((M, N), cg, 'tensorflow') y = session.run(f(ma, mb)) c = session.run(tf.einsum("ij,kl", ma, mb)) assert (y == c).all() cg = CodegenArrayElementwiseAdd(M, N) assert tensorflow_code(cg) == 'tensorflow.math.add(M, N)' f = lambdify((M, N), cg, 'tensorflow') y = session.run(f(ma, mb)) c = session.run(ma + mb) assert (y == c).all() cg = CodegenArrayElementwiseAdd(M, N, P) assert tensorflow_code(cg) == \ 'tensorflow.math.add(tensorflow.math.add(M, N), P)' f = lambdify((M, N, P), cg, 'tensorflow') y = session.run(f(ma, mb, mc)) c = session.run(ma + mb + mc) assert (y == c).all() cg = CodegenArrayElementwiseAdd(M, N, P, Q) assert tensorflow_code(cg) == \ 'tensorflow.math.add(' \ 'tensorflow.math.add(tensorflow.math.add(M, N), P), Q)' f = lambdify((M, N, P, Q), cg, 'tensorflow') y = session.run(f(ma, mb, mc, md)) c = session.run(ma + mb + mc + md) assert (y == c).all() cg = CodegenArrayPermuteDims(M, [1, 0]) assert tensorflow_code(cg) == 'tensorflow.transpose(M, [1, 0])' f = lambdify((M, ), cg, 'tensorflow') y = session.run(f(ma)) c = session.run(tf.transpose(ma)) assert (y == c).all() cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 2, 3, 0]) assert tensorflow_code(cg) == \ 'tensorflow.transpose(' \ 'tensorflow.linalg.einsum("ab,cd", M, N), [1, 2, 3, 0])' f = lambdify((M, N), cg, 'tensorflow') y = session.run(f(ma, mb)) c = session.run(tf.transpose(tf.einsum("ab,cd", ma, mb), [1, 2, 3, 0])) assert (y == c).all() cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2)) assert tensorflow_code(cg) == \ 'tensorflow.linalg.einsum("ab,bc->acb", M, N)' f = lambdify((M, N), cg, 'tensorflow') y = session.run(f(ma, mb)) c = session.run(tf.einsum("ab,bc->acb", ma, mb)) assert (y == c).all()
def test_array_wrong_permutation_size(): cg = CodegenArrayTensorProduct(M, N) raises(ValueError, lambda: CodegenArrayPermuteDims(cg, [1, 0])) raises(ValueError, lambda: CodegenArrayPermuteDims(cg, [1, 0, 2, 3, 5, 4]))
def test_parsing_of_matrix_expressions(): expr = M*N assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)) expr = Transpose(M) assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0]) expr = M*Transpose(N) assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2)) expr = 3*M*N res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr == rexpr expr = 3*M + N*M.T*M + 4*k*N res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr == rexpr expr = Inverse(M)*N rexpr = recognize_matrix_expression(parse_matrix_expression(expr)) assert expr == rexpr expr = M**2 rexpr = recognize_matrix_expression(parse_matrix_expression(expr)) assert expr == rexpr expr = M*(2*N + 3*M) res = parse_matrix_expression(expr) rexpr = recognize_matrix_expression(res) assert expr == rexpr expr = Trace(M) result = CodegenArrayContraction(M, (0, 1)) assert parse_matrix_expression(expr) == result expr = 3*Trace(M) result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M), (0, 1)) assert parse_matrix_expression(expr) == result expr = 3*Trace(Trace(M) * M) result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3)) assert parse_matrix_expression(expr) == result expr = 3*Trace(M)**2 result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3)) assert parse_matrix_expression(expr) == result expr = HadamardProduct(M, N) result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3)) assert parse_matrix_expression(expr) == result expr = HadamardPower(M, 2) result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, M), (0, 2), (1, 3)) assert parse_matrix_expression(expr) == result expr = M**2 assert isinstance(expr, MatPow) assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, M), (1, 2))