def test_arrayexpr_permutedims_sink(): cg = _permute_dims(_array_tensor_product(M, N), [0, 1, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_tensor_product(M, _permute_dims(N, [1, 0])) cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0])) cg = _permute_dims(_array_tensor_product(M, N), [3, 2, 1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_tensor_product(_permute_dims(N, [1, 0]), _permute_dims(M, [1, 0])) cg = _permute_dims(_array_contraction(_array_tensor_product(M, N), (1, 2)), [1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N), [[0, 3]]), (1, 2)) cg = _permute_dims(_array_tensor_product(M, N), [1, 0, 3, 2], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_tensor_product(_permute_dims(M, [1, 0]), _permute_dims(N, [1, 0])) cg = _permute_dims(_array_contraction(_array_tensor_product(M, N, P), (1, 2), (3, 4)), [1, 0], nest_permutation=False) sunk = nest_permutation(cg) assert sunk == _array_contraction(_permute_dims(_array_tensor_product(M, N, P), [[0, 5]]), (1, 2), (3, 4))
def test_arrayexpr_convert_indexed_to_array_expression(): s = Sum(A[i] * B[i], (i, 0, 3)) cg = convert_indexed_to_array(s) assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1)) expr = M * N result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2)) elem = expr[i, j] assert convert_indexed_to_array(elem) == result expr = M * N * M elem = expr[i, j] result = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5)) cg = convert_indexed_to_array(elem) assert cg == result cg = convert_indexed_to_array((M * N * P)[i, j]) assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)) cg = convert_indexed_to_array((M * N.T * P)[i, j]) assert cg == _array_contraction(ArrayTensorProduct(M, N, P), (1, 3), (2, 4)) expr = -2 * M * N elem = expr[i, j] cg = convert_indexed_to_array(elem) assert cg == ArrayContraction(ArrayTensorProduct(-2, M, N), (1, 2))
def test_arrayexpr_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 _array_tensor_product(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n) assert _array_tensor_product(M, N, zm1) == ZeroArray(k, k, k, k, m, n) assert _array_contraction(za1, (3,)) == ZeroArray(k, l, m) assert _array_contraction(zm1, (1,)) == ZeroArray(m) assert _array_contraction(za2, (1, 2)) == ZeroArray(k, n) assert _array_contraction(zm2, (0, 1)) == 0 assert _array_diagonal(za2, (1, 2)) == ZeroArray(k, n, m) assert _array_diagonal(zm2, (0, 1)) == ZeroArray(m) assert _permute_dims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k) assert _permute_dims(zm1, [1, 0]) == ZeroArray(n, m) assert _array_add(za1) == za1 assert _array_add(zm1) == ZeroArray(m, n) tp1 = _array_tensor_product(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n)) assert _array_add(tp1, za1) == tp1 tp2 = _array_tensor_product(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n)) assert _array_add(tp1, za1, tp2) == _array_add(tp1, tp2) assert _array_add(M, zm3) == M assert _array_add(M, N, zm3) == _array_add(M, N)
def test_arrayexpr_split_multiple_contractions(): 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) cg = _array_contraction( _array_tensor_product(A.T, a, b, b.T, (A * X * b).applyfunc(cos)), (1, 2, 8), (5, 6, 9)) expected = _array_contraction( _array_tensor_product(A.T, DiagMatrix(a), OneArray(1), b, b.T, (A * X * b).applyfunc(cos)), (1, 3), (2, 9), (6, 7, 10)) assert cg.split_multiple_contractions().dummy_eq(expected) # Check no overlap of lines: cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8), (3, 7)) assert cg.split_multiple_contractions() == cg cg = _array_contraction(_array_tensor_product(a, b, A), (0, 2, 4), (1, 3)) assert cg.split_multiple_contractions() == cg
def test_arrayexpr_convert_array_to_matrix(): cg = _array_contraction(_array_tensor_product(M), (0, 1)) assert convert_array_to_matrix(cg) == Trace(M) cg = _array_contraction(_array_tensor_product(M, N), (0, 1), (2, 3)) assert convert_array_to_matrix(cg) == Trace(M) * Trace(N) cg = _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2)) assert convert_array_to_matrix(cg) == Trace(M * N) cg = _array_contraction(_array_tensor_product(M, N), (0, 2), (1, 3)) assert convert_array_to_matrix(cg) == Trace(M * N.T) cg = convert_matrix_to_array(M * N * P) assert convert_array_to_matrix(cg) == M * N * P cg = convert_matrix_to_array(M * N.T * P) assert convert_array_to_matrix(cg) == M * N.T * P cg = _array_contraction(_array_tensor_product(M,N,P,Q), (1, 2), (5, 6)) assert convert_array_to_matrix(cg) == _array_tensor_product(M * N, P * Q) cg = _array_contraction(_array_tensor_product(-2, M, N), (1, 2)) assert convert_array_to_matrix(cg) == -2 * M * N a = MatrixSymbol("a", k, 1) b = MatrixSymbol("b", k, 1) c = MatrixSymbol("c", k, 1) cg = PermuteDims( _array_contraction( _array_tensor_product( a, ArrayAdd( _array_tensor_product(b, c), _array_tensor_product(c, b), ) ), (2, 4)), [0, 1, 3, 2]) assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b) za = ZeroArray(m, n) assert convert_array_to_matrix(za) == ZeroMatrix(m, n) cg = _array_tensor_product(3, M) assert convert_array_to_matrix(cg) == 3 * M # Partial conversion to matrix multiplication: expr = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 4, 6)) assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(M.T*N, P, Q), (0, 2, 4)) x = MatrixSymbol("x", k, 1) cg = PermuteDims( _array_contraction(_array_tensor_product(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))), (0, 5)), Permutation(1, 2, 3)) assert convert_array_to_matrix(cg) == x expr = ArrayAdd(M, PermuteDims(M, [1, 0])) assert convert_array_to_matrix(expr) == M + Transpose(M)
def test_arrayexpr_contraction_construction(): cg = _array_contraction(A) assert cg == A cg = _array_contraction(_array_tensor_product(A, B), (1, 0)) assert cg == _array_contraction(_array_tensor_product(A, B), (0, 1)) cg = _array_contraction(_array_tensor_product(M, N), (0, 1)) indtup = cg._get_contraction_tuples() assert indtup == [[(0, 0), (0, 1)]] assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(0, 1)] cg = _array_contraction(_array_tensor_product(M, N), (1, 2)) indtup = cg._get_contraction_tuples() assert indtup == [[(0, 1), (1, 0)]] assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(1, 2)] cg = _array_contraction(_array_tensor_product(M, M, N), (1, 4), (2, 5)) indtup = cg._get_contraction_tuples() assert indtup == [[(0, 0), (1, 1)], [(0, 1), (2, 0)]] assert cg._contraction_tuples_to_contraction_indices(cg.expr, indtup) == [(0, 3), (1, 4)] # Test removal of trivial contraction: assert _array_contraction(a, (1,)) == a assert _array_contraction( _array_tensor_product(a, b), (0, 2), (1,), (3,)) == _array_contraction( _array_tensor_product(a, b), (0, 2))
def test_array_contraction_to_diagonal_multiple_identities(): expr = _array_contraction(_array_tensor_product(A, B, I, C), (1, 2, 4), (5, 6)) assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, []) assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(A, B, C), (1, 2, 4)) expr = _array_contraction(_array_tensor_product(A, I, I), (1, 2, 4)) assert _array_contraction_to_diagonal_multiple_identity(expr) == (A, [2]) assert convert_array_to_matrix(expr) == A expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 4), (3, 6)) assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, []) expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 3, 4, 6)) assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
def test_arrayexpr_nested_array_elementwise_add(): cg = _array_contraction( _array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)), (1, 2)) result = _array_add( _array_contraction(_array_tensor_product(M, N), (1, 2)), _array_contraction(_array_tensor_product(N, M), (1, 2))) assert cg == result cg = _array_diagonal( _array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)), (1, 2)) result = _array_add(_array_diagonal(_array_tensor_product(M, N), (1, 2)), _array_diagonal(_array_tensor_product(N, M), (1, 2))) assert cg == result
def _(expr: Inverse, x: Expr): mat = expr.I dexpr = array_derive(mat, x) tp = _array_tensor_product(-expr, dexpr, expr) mp = _array_contraction(tp, (1, 4), (5, 6)) pp = _permute_dims(mp, [1, 2, 0, 3]) return pp
def _(expr: ArrayContraction): new_expr, removed0 = _array_contraction_to_diagonal_multiple_identity(expr) if new_expr != expr: new_expr2, removed1 = _remove_trivial_dims(_array2matrix(new_expr)) removed = _combine_removed(-1, removed0, removed1) return new_expr2, removed rank1 = get_rank(expr) expr, removed1 = remove_identity_matrices(expr) if not isinstance(expr, ArrayContraction): expr2, removed2 = _remove_trivial_dims(expr) return expr2, _combine_removed(rank1, removed1, removed2) newexpr, removed2 = _remove_trivial_dims(expr.expr) shifts = list( accumulate( [1 if i in removed2 else 0 for i in range(get_rank(expr.expr))])) new_contraction_indices = [ tuple(j for j in i if j not in removed2) for i in expr.contraction_indices ] # Remove possible empty tuples "()": new_contraction_indices = [ i for i in new_contraction_indices if len(i) > 0 ] contraction_indices_flat = [j for i in expr.contraction_indices for j in i] removed2 = [i for i in removed2 if i not in contraction_indices_flat] new_contraction_indices = [ tuple(j - shifts[j] for j in i) for i in new_contraction_indices ] # Shift removed2: removed2 = ArrayContraction._push_indices_up(expr.contraction_indices, removed2) removed = _combine_removed(rank1, removed1, removed2) return _array_contraction(newexpr, *new_contraction_indices), list(removed)
def _(expr: ArrayContraction, x: Expr): fd = array_derive(expr.expr, x) rank_x = len(get_shape(x)) contraction_indices = expr.contraction_indices new_contraction_indices = [ tuple(j + rank_x for j in i) for i in contraction_indices ] return _array_contraction(fd, *new_contraction_indices)
def _a2m_mul(*args): if not any(isinstance(i, _CodegenArrayAbstract) for i in args): from sympy.matrices.expressions.matmul import MatMul return MatMul(*args).doit() else: return _array_contraction( _array_tensor_product(*args), *[(2 * i - 1, 2 * i) for i in range(1, len(args))])
def _support_function_tp1_recognize(contraction_indices, args): if len(contraction_indices) == 0: return _a2m_tensor_product(*args) ac = _array_contraction(_array_tensor_product(*args), *contraction_indices) editor = _EditArrayContraction(ac) editor.track_permutation_start() while True: flag_stop: bool = True for i, arg_with_ind in enumerate(editor.args_with_ind): if not isinstance(arg_with_ind.element, MatrixExpr): continue first_index = arg_with_ind.indices[0] second_index = arg_with_ind.indices[1] first_frequency = editor.count_args_with_index(first_index) second_frequency = editor.count_args_with_index(second_index) if first_index is not None and first_frequency == 1 and first_index == second_index: flag_stop = False arg_with_ind.element = Trace(arg_with_ind.element)._normalize() arg_with_ind.indices = [] break scan_indices = [] if first_frequency == 2: scan_indices.append(first_index) if second_frequency == 2: scan_indices.append(second_index) candidate, transpose, found_index = _get_candidate_for_matmul_from_contraction( scan_indices, editor.args_with_ind[i + 1:]) if candidate is not None: flag_stop = False editor.track_permutation_merge(arg_with_ind, candidate) transpose1 = found_index == first_index new_arge, other_index = _insert_candidate_into_editor( editor, arg_with_ind, candidate, transpose1, transpose) if found_index == first_index: new_arge.indices = [second_index, other_index] else: new_arge.indices = [first_index, other_index] set_indices = set(new_arge.indices) if len(set_indices) == 1 and set_indices != {None}: # This is a trace: new_arge.element = Trace(new_arge.element)._normalize() new_arge.indices = [] editor.args_with_ind[i] = new_arge # TODO: is this break necessary? break if flag_stop: break editor.refresh_indices() return editor.to_array_contraction()
def test_arrayexpr_convert_array_contraction_tp_additions(): a = ArrayAdd( _array_tensor_product(M, N), _array_tensor_product(N, M) ) tp = _array_tensor_product(P, a, Q) expr = _array_contraction(tp, (3, 4)) expected = _array_tensor_product( P, ArrayAdd( _array_contraction(_array_tensor_product(M, N), (1, 2)), _array_contraction(_array_tensor_product(N, M), (1, 2)), ), Q ) assert expr == expected assert convert_array_to_matrix(expr) == _array_tensor_product(P, M * N + N * M, Q) expr = _array_contraction(tp, (1, 2), (3, 4), (5, 6)) result = _array_contraction( _array_tensor_product( P, ArrayAdd( _array_contraction(_array_tensor_product(M, N), (1, 2)), _array_contraction(_array_tensor_product(N, M), (1, 2)), ), Q ), (1, 2), (3, 4)) assert expr == result assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
def _array_diag2contr_diagmatrix(expr: ArrayDiagonal): if isinstance(expr.expr, ArrayTensorProduct): args = list(expr.expr.args) diag_indices = list(expr.diagonal_indices) mapping = _get_mapping_from_subranks( [_get_subrank(arg) for arg in args]) tuple_links = [[mapping[j] for j in i] for i in diag_indices] contr_indices = [] total_rank = get_rank(expr) replaced = [False for arg in args] for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)): if len(abs_pos) != 2: continue (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos arg1 = args[pos1_outer] arg2 = args[pos2_outer] if get_rank(arg1) != 2 or get_rank(arg2) != 2: if replaced[pos1_outer]: diag_indices[i] = None if replaced[pos2_outer]: diag_indices[i] = None continue pos1_in2 = 1 - pos1_inner pos2_in2 = 1 - pos2_inner if arg1.shape[pos1_in2] == 1: if arg1.shape[pos1_inner] != 1: darg1 = DiagMatrix(arg1) else: darg1 = arg1 args.append(darg1) contr_indices.append( ((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner))) total_rank += 1 diag_indices[i] = None args[pos1_outer] = OneArray(arg1.shape[pos1_in2]) replaced[pos1_outer] = True elif arg2.shape[pos2_in2] == 1: if arg2.shape[pos2_inner] != 1: darg2 = DiagMatrix(arg2) else: darg2 = arg2 args.append(darg2) contr_indices.append( ((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner))) total_rank += 1 diag_indices[i] = None args[pos2_outer] = OneArray(arg2.shape[pos2_in2]) replaced[pos2_outer] = True diag_indices_new = [i for i in diag_indices if i is not None] cumul = list(accumulate([0] + [get_rank(arg) for arg in args])) contr_indices2 = [ tuple(cumul[a] + b for a, b in i) for i in contr_indices ] tc = _array_contraction(_array_tensor_product(*args), *contr_indices2) td = _array_diagonal(tc, *diag_indices_new) return td return expr
def test_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix(): cg = _array_diagonal(_array_tensor_product(M, a), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == _array_contraction( _array_tensor_product(M, OneArray(1), DiagMatrix(a)), (1, 3)) raises(ValueError, lambda: _array_diagonal(_array_tensor_product(a, M), (1, 2))) cg = _array_diagonal(_array_tensor_product(a.T, M), (1, 2)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == _array_contraction( _array_tensor_product(OneArray(1), M, DiagMatrix(a.T)), (1, 4)) cg = _array_diagonal(_array_tensor_product(a.T, M, N, b.T), (1, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == _array_contraction( _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9)) cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 2), (4, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == _array_contraction( _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9)) cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 4), (3, 7)) res = _array_diag2contr_diagmatrix(cg) assert res.shape == cg.shape assert res == _array_contraction( _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9)) I1 = Identity(1) x = MatrixSymbol("x", k, 1) A = MatrixSymbol("A", k, k) cg = _array_diagonal(_array_tensor_product(x, A.T, I1), (0, 2)) assert _array_diag2contr_diagmatrix(cg).shape == cg.shape assert _array2matrix(cg).shape == cg.shape
def test_arrayexpr_contraction_permutation_mix(): Me = M.subs(k, 3).as_explicit() Ne = N.subs(k, 3).as_explicit() cg1 = _array_contraction(PermuteDims(_array_tensor_product(M, N), Permutation([0, 2, 1, 3])), (2, 3)) cg2 = _array_contraction(_array_tensor_product(M, N), (1, 3)) assert cg1 == cg2 cge1 = tensorcontraction(permutedims(tensorproduct(Me, Ne), Permutation([0, 2, 1, 3])), (2, 3)) cge2 = tensorcontraction(tensorproduct(Me, Ne), (1, 3)) assert cge1 == cge2 cg1 = _permute_dims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2])) cg2 = _array_tensor_product(M, _permute_dims(N, Permutation([1, 0]))) assert cg1 == cg2 cg1 = _array_contraction( _permute_dims( _array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5) ) cg2 = _array_contraction( _array_tensor_product(M, N, P, _permute_dims(Q, Permutation([1, 0]))), (1, 5), (2, 3) ) assert cg1 == cg2 cg1 = _array_contraction( _permute_dims( _array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 2, 7, 5, 3])), (0, 1), (2, 6), (3, 7) ) cg2 = _permute_dims( _array_contraction( _array_tensor_product(M, P, Q, N), (0, 1), (2, 3), (4, 7)), [1, 0] ) assert cg1 == cg2 cg1 = _array_contraction( _permute_dims( _array_tensor_product(M, N, P, Q), Permutation([1, 0, 4, 6, 7, 2, 5, 3])), (0, 1), (2, 6), (3, 7) ) cg2 = _permute_dims( _array_contraction( _array_tensor_product(_permute_dims(M, [1, 0]), N, P, Q), (0, 1), (3, 6), (4, 5) ), Permutation([1, 0]) ) assert cg1 == cg2
def test_convert_array_element_to_matrix(): expr = ArrayElement(M, (i, j)) assert convert_array_to_matrix(expr) == MatrixElement(M, i, j) expr = ArrayElement(_array_contraction(_array_tensor_product(M, N), (1, 3)), (i, j)) assert convert_array_to_matrix(expr) == MatrixElement(M*N.T, i, j) expr = ArrayElement(_array_tensor_product(M, N), (i, j, m, n)) assert convert_array_to_matrix(expr) == expr
def _(expr: ArrayContraction): expr = expr.flatten_contraction_of_diagonal() expr = identify_removable_identity_matrices(expr) expr = expr.split_multiple_contractions() expr = identify_hadamard_products(expr) if not isinstance(expr, ArrayContraction): return _array2matrix(expr) subexpr = expr.expr contraction_indices: tTuple[tTuple[int]] = expr.contraction_indices if contraction_indices == ((0, ), (1, )) or (contraction_indices == ( (0, ), ) and subexpr.shape[1] == 1) or (contraction_indices == ( (1, ), ) and subexpr.shape[0] == 1): shape = subexpr.shape subexpr = _array2matrix(subexpr) if isinstance(subexpr, MatrixExpr): return OneMatrix(1, shape[0]) * subexpr * OneMatrix(shape[1], 1) if isinstance(subexpr, ArrayTensorProduct): newexpr = _array_contraction(_array2matrix(subexpr), *contraction_indices) contraction_indices = newexpr.contraction_indices if any(i > 2 for i in newexpr.subranks): addends = _array_add(*[ _a2m_tensor_product(*j) for j in itertools.product(*[ i.args if isinstance(i, ArrayAdd) else [i] for i in expr.expr.args ]) ]) newexpr = _array_contraction(addends, *contraction_indices) if isinstance(newexpr, ArrayAdd): ret = _array2matrix(newexpr) return ret assert isinstance(newexpr, ArrayContraction) ret = _support_function_tp1_recognize(contraction_indices, list(newexpr.expr.args)) return ret elif not isinstance(subexpr, _CodegenArrayAbstract): ret = _array2matrix(subexpr) if isinstance(ret, MatrixExpr): assert expr.contraction_indices == ((0, 1), ) return _a2m_trace(ret) else: return _array_contraction(ret, *expr.contraction_indices)
def test_edit_array_contraction(): cg = _array_contraction(_array_tensor_product(A, B, C, D), (1, 2, 5)) ecg = _EditArrayContraction(cg) assert ecg.to_array_contraction() == cg ecg.args_with_ind[1], ecg.args_with_ind[2] = ecg.args_with_ind[2], ecg.args_with_ind[1] assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, C, B, D), (1, 3, 4)) ci = ecg.get_new_contraction_index() new_arg = _ArgE(X) new_arg.indices = [ci, ci] ecg.args_with_ind.insert(2, new_arg) assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, C, X, B, D), (1, 3, 6), (4, 5)) assert ecg.get_contraction_indices() == [[1, 3, 6], [4, 5]] assert [[tuple(j) for j in i] for i in ecg.get_contraction_indices_to_ind_rel_pos()] == [[(0, 1), (1, 1), (3, 0)], [(2, 0), (2, 1)]] assert [list(i) for i in ecg.get_mapping_for_index(0)] == [[0, 1], [1, 1], [3, 0]] assert [list(i) for i in ecg.get_mapping_for_index(1)] == [[2, 0], [2, 1]] raises(ValueError, lambda: ecg.get_mapping_for_index(2)) ecg.args_with_ind.pop(1) assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, B, D), (1, 4), (2, 3)) ecg.args_with_ind[0].indices[1] = ecg.args_with_ind[1].indices[0] ecg.args_with_ind[1].indices[1] = ecg.args_with_ind[2].indices[0] assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, B, D), (1, 2), (3, 4)) ecg.insert_after(ecg.args_with_ind[1], _ArgE(C)) assert ecg.to_array_contraction() == _array_contraction(_array_tensor_product(A, X, C, B, D), (1, 2), (3, 6))
def test_arrayexpr_array_shape(): expr = _array_tensor_product(M, N, P, Q) assert expr.shape == (k, k, k, k, k, k, k, k) Z = MatrixSymbol("Z", m, n) expr = _array_tensor_product(M, Z) assert expr.shape == (k, k, m, n) expr2 = _array_contraction(expr, (0, 1)) assert expr2.shape == (m, n) expr2 = _array_diagonal(expr, (0, 1)) assert expr2.shape == (m, n, k) exprp = _permute_dims(expr, [2, 1, 3, 0]) assert exprp.shape == (m, k, n, k) expr3 = _array_tensor_product(N, Z) expr2 = _array_add(expr, expr3) assert expr2.shape == (k, k, m, n) # Contraction along axes with discordant dimensions: raises(ValueError, lambda: _array_contraction(expr, (1, 2))) # Also diagonal needs the same dimensions: raises(ValueError, lambda: _array_diagonal(expr, (1, 2))) # Diagonal requires at least to axes to compute the diagonal: raises(ValueError, lambda: _array_diagonal(expr, (1,)))
def test_arrayexpr_convert_index_to_array_support_function(): expr = M[i, j] assert _convert_indexed_to_array(expr) == (M, (i, j)) expr = M[i, j] * N[k, l] assert _convert_indexed_to_array(expr) == (ArrayTensorProduct(M, N), (i, j, k, l)) expr = M[i, j] * N[j, k] assert _convert_indexed_to_array(expr) == (ArrayDiagonal( ArrayTensorProduct(M, N), (1, 2)), (i, k, j)) expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1)) assert _convert_indexed_to_array(expr) == (ArrayContraction( ArrayTensorProduct(M, N), (1, 2)), (i, k)) expr = M[i, j] + N[i, j] assert _convert_indexed_to_array(expr) == (ArrayAdd(M, N), (i, j)) expr = M[i, j] + N[j, i] assert _convert_indexed_to_array(expr) == (ArrayAdd( M, PermuteDims(N, Permutation([1, 0]))), (i, j)) expr = M[i, j] + M[j, i] assert _convert_indexed_to_array(expr) == (ArrayAdd( M, PermuteDims(M, Permutation([1, 0]))), (i, j)) expr = (M * N * P)[i, j] assert _convert_indexed_to_array(expr) == (_array_contraction( ArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j)) expr = expr.function # Disregard summation in previous expression ret1, ret2 = _convert_indexed_to_array(expr) assert ret1 == ArrayDiagonal(ArrayTensorProduct(M, N, P), (1, 2), (3, 4)) assert str(ret2) == "(i, j, _i_1, _i_2)" expr = KroneckerDelta(i, j) * M[i, k] assert _convert_indexed_to_array(expr) == (M, ({i, j}, k)) expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * M[i, l] assert _convert_indexed_to_array(expr) == (M, ({i, j, k}, l)) expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l]) assert _convert_indexed_to_array(expr) == (_array_diagonal( _array_add( ArrayTensorProduct(M, N), _permute_dims(ArrayTensorProduct(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 _convert_indexed_to_array(expr) == (_array_diagonal( _array_add( ArrayTensorProduct(M, N), _permute_dims(ArrayTensorProduct(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 _convert_indexed_to_array(expr) == (M, ({i, j, k, m, n}, 0)) expr = M[i, i] assert _convert_indexed_to_array(expr) == (ArrayDiagonal(M, (0, 1)), (i, ))
def test_arrayexpr_canonicalize_diagonal_contraction(): tp = _array_tensor_product(M, N, P, Q) expr = _array_contraction(_array_diagonal(tp, (1, 3, 4)), (0, 3)) result = _array_diagonal(_array_contraction(_array_tensor_product(M, N, P, Q), (0, 6)), (0, 2, 3)) assert expr == result expr = _array_contraction(_array_diagonal(tp, (0, 1, 2, 3, 7)), (1, 2, 3)) result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 2, 3, 5, 6, 7)) assert expr == result expr = _array_contraction(_array_diagonal(tp, (0, 2, 6, 7)), (1, 2, 3)) result = _array_diagonal(_array_contraction(tp, (3, 4, 5)), (0, 2, 3, 4)) assert expr == result td = _array_diagonal(_array_tensor_product(M, N, P, Q), (0, 3)) expr = _array_contraction(td, (2, 1), (0, 4, 6, 5, 3)) result = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 1, 3, 5, 6, 7), (2, 4)) assert expr == result
def test_arrayexpr_convert_array_to_implicit_matmul(): # Trivial dimensions are suppressed, so the result can be expressed in matrix form: cg = _array_tensor_product(a, b) assert convert_array_to_matrix(cg) == a * b.T cg = _array_tensor_product(a, b, I) assert convert_array_to_matrix(cg) == _array_tensor_product(a*b.T, I) cg = _array_tensor_product(I, a, b) assert convert_array_to_matrix(cg) == _array_tensor_product(I, a*b.T) cg = _array_tensor_product(a, I, b) assert convert_array_to_matrix(cg) == _array_tensor_product(a, I, b) cg = _array_contraction(_array_tensor_product(I, I), (1, 2)) assert convert_array_to_matrix(cg) == I cg = PermuteDims(_array_tensor_product(I, Identity(1)), [0, 2, 1, 3]) assert convert_array_to_matrix(cg) == I
def test_identify_removable_identity_matrices(): D = DiagonalMatrix(MatrixSymbol("D", k, k)) cg = _array_contraction(_array_tensor_product(A, B, I), (1, 2, 4, 5)) expected = _array_contraction(_array_tensor_product(A, B), (1, 2)) assert identify_removable_identity_matrices(cg) == expected cg = _array_contraction(_array_tensor_product(A, B, C, I), (1, 3, 5, 6, 7)) expected = _array_contraction(_array_tensor_product(A, B, C), (1, 3, 5)) assert identify_removable_identity_matrices(cg) == expected # Tests with diagonal matrices: cg = _array_contraction(_array_tensor_product(A, B, D), (1, 2, 4, 5)) ret = identify_removable_identity_matrices(cg) expected = _array_contraction(_array_tensor_product(A, B, D), (1, 4), (2, 5)) assert ret == expected cg = _array_contraction(_array_tensor_product(A, B, D, M, N), (1, 2, 4, 5, 6, 8)) ret = identify_removable_identity_matrices(cg) assert ret == cg
def test_arrayexpr_permute_tensor_product(): cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 1, 0, 5, 4, 6, 7])) cg2 = _array_tensor_product(N, _permute_dims(M, [1, 0]), _permute_dims(P, [1, 0]), Q) assert cg1 == cg2 # TODO: reverse operation starting with `PermuteDims` and getting down to `bb`... cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 5, 0, 1, 6, 7])) cg2 = _array_tensor_product(N, P, M, Q) assert cg1 == cg2 cg1 = _permute_dims(_array_tensor_product(M, N, P, Q), Permutation([2, 3, 4, 6, 5, 7, 0, 1])) assert cg1.expr == _array_tensor_product(N, P, Q, M) assert cg1.permutation == Permutation([0, 1, 2, 4, 3, 5, 6, 7]) cg1 = _array_contraction( _permute_dims( _array_tensor_product(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]) cg2 = _permute_dims(_array_contraction(_array_tensor_product(Q, Q, N, M), (3, 5, 6)), [0, 2, 3, 1, 4]) assert cg1 == cg2 cg1 = _array_contraction( _array_contraction( _array_contraction( _array_contraction( _permute_dims( _array_tensor_product(N, Q, Q, M), [2, 1, 5, 4, 0, 3, 6, 7]), [1, 2, 6]), [1, 3, 4]), [1]), [0]) cg2 = _array_contraction(_array_tensor_product(M, N, Q, Q), (0, 3, 5), (1, 4, 7), (2,), (6,)) assert cg1 == cg2
def test_arrayexpr_convert_array_to_matrix2(): cg = _array_contraction(_array_tensor_product(M, N), (1, 3)) assert convert_array_to_matrix(cg) == M * N.T cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2])) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0]))) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_contraction( PermuteDims( _array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])), (1, 2), (3, 5) ) assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T) cg = _array_contraction( _array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))), (1, 5), (2, 3) ) assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T) cg = _array_tensor_product(M, PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T) cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T) cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0])) assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T) cg = _array_contraction(M, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1) cg = _array_contraction(x, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x Xm = MatrixSymbol("Xm", m, n) cg = _array_contraction(Xm, (0,), (1,)) assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims(): # Tensor Product: assert _remove_trivial_dims(_array_tensor_product(a, b)) == (a * b.T, [1, 3]) assert _remove_trivial_dims(_array_tensor_product(a.T, b)) == (a * b.T, [0, 3]) assert _remove_trivial_dims(_array_tensor_product(a, b.T)) == (a * b.T, [1, 2]) assert _remove_trivial_dims(_array_tensor_product(a.T, b.T)) == (a * b.T, [0, 2]) assert _remove_trivial_dims(_array_tensor_product(I, a.T, b.T)) == (_array_tensor_product(I, a * b.T), [2, 4]) assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T)) == (_array_tensor_product(a.T, I, b.T), []) assert _remove_trivial_dims(_array_tensor_product(a, I)) == (_array_tensor_product(a, I), []) assert _remove_trivial_dims(_array_tensor_product(I, a)) == (_array_tensor_product(I, a), []) assert _remove_trivial_dims(_array_tensor_product(a.T, b.T, c, d)) == ( _array_tensor_product(a * b.T, c * d.T), [0, 2, 5, 7]) assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T, c, d, I)) == ( _array_tensor_product(a.T, I, b*c.T, d, I), [4, 7]) # Addition: cg = ArrayAdd(_array_tensor_product(a, b), _array_tensor_product(c, d)) assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3]) # Permute Dims: cg = PermuteDims(_array_tensor_product(a, b), Permutation(3)(1, 2)) assert _remove_trivial_dims(cg) == (a * b.T, [2, 3]) cg = PermuteDims(_array_tensor_product(a, I, b), Permutation(5)(1, 2, 3, 4)) assert _remove_trivial_dims(cg) == (cg, []) cg = PermuteDims(_array_tensor_product(I, b, a), Permutation(5)(1, 2, 4, 5, 3)) assert _remove_trivial_dims(cg) == (PermuteDims(_array_tensor_product(I, b * a.T), [0, 2, 3, 1]), [4, 5]) # Diagonal: cg = _array_diagonal(_array_tensor_product(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, []) # Contraction: cg = _array_contraction(_array_tensor_product(M, a), (1, 2)) assert _remove_trivial_dims(cg) == (cg, []) # A few more cases to test the removal and shift of nested removed axes # with array contractions and array diagonals: tp = _array_tensor_product( OneMatrix(1, 1), M, x, OneMatrix(1, 1), Identity(1), ) expr = _array_contraction(tp, (1, 8)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 5, 6, 7] expr = _array_contraction(tp, (1, 8), (3, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 3, 4, 5] expr = _array_diagonal(tp, (1, 8)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 5, 6, 7, 8] expr = _array_diagonal(tp, (1, 8), (3, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [0, 3, 4, 5, 6] expr = _array_diagonal(_array_contraction(_array_tensor_product(A, x, I, I1), (1, 2, 5)), (1, 4)) rexpr, removed = _remove_trivial_dims(expr) assert removed == [2, 3] cg = _array_diagonal(_array_tensor_product(PermuteDims(_array_tensor_product(x, I1), Permutation(1, 2, 3)), (x.T*x).applyfunc(sqrt)), (2, 4), (3, 5)) rexpr, removed = _remove_trivial_dims(cg) assert removed == [1, 2] # Contractions with identity matrices need to be followed by a permutation # in order cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8)) ret, removed = _remove_trivial_dims(cg) assert ret == PermuteDims(_array_tensor_product(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1]) assert removed == [] cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8), (3, 4)) ret, removed = _remove_trivial_dims(cg) assert ret == PermuteDims(_array_contraction(_array_tensor_product(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1]) assert removed == [] # Trivial matrices are sometimes inserted into MatMul expressions: cg = _array_tensor_product(b*b.T, a.T*a) ret, removed = _remove_trivial_dims(cg) assert ret == b*a.T*a*b.T assert removed == [2, 3] Xs = ArraySymbol("X", (3, 2, k)) cg = _array_tensor_product(M, Xs, b.T*c, a*a.T, b*b.T, c.T*d) ret, removed = _remove_trivial_dims(cg) assert ret == _array_tensor_product(M, Xs, a*b.T*c*c.T*d*a.T, b*b.T) assert removed == [5, 6, 11, 12] cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5)) assert _remove_trivial_dims(cg) == (PermuteDims(_array_diagonal(_array_tensor_product(I, x), (1, 2)), Permutation(1, 2)), [1]) expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2)) assert _remove_trivial_dims(expr) == (PermuteDims(_array_tensor_product(DiagMatrix(x), y), [1, 2, 3, 0]), [0]) expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2), (3, 4)) assert _remove_trivial_dims(expr) == (expr, [])
def test_array_expressions_no_canonicalization(): tp = _array_tensor_product(M, N, P) # ArrayTensorProduct: expr = ArrayTensorProduct(tp, N) assert str(expr) == "ArrayTensorProduct(ArrayTensorProduct(M, N, P), N)" assert expr.doit() == ArrayTensorProduct(M, N, P, N) expr = ArrayTensorProduct(ArrayContraction(M, (0, 1)), N) assert str(expr) == "ArrayTensorProduct(ArrayContraction(M, (0, 1)), N)" assert expr.doit() == ArrayContraction(ArrayTensorProduct(M, N), (0, 1)) expr = ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N) assert str(expr) == "ArrayTensorProduct(ArrayDiagonal(M, (0, 1)), N)" assert expr.doit() == PermuteDims(ArrayDiagonal(ArrayTensorProduct(M, N), (0, 1)), [2, 0, 1]) expr = ArrayTensorProduct(PermuteDims(M, [1, 0]), N) assert str(expr) == "ArrayTensorProduct(PermuteDims(M, (0 1)), N)" assert expr.doit() == PermuteDims(ArrayTensorProduct(M, N), [1, 0, 2, 3]) # ArrayContraction: expr = ArrayContraction(_array_contraction(tp, (0, 2)), (0, 1)) assert isinstance(expr, ArrayContraction) assert isinstance(expr.expr, ArrayContraction) assert str(expr) == "ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))" assert expr.doit() == ArrayContraction(tp, (0, 2), (1, 3)) expr = ArrayContraction(ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1)), (0, 1)) assert expr.doit() == ArrayContraction(tp, (0, 1), (2, 3), (4, 5)) # assert expr._canonicalize() == ArrayContraction(ArrayContraction(tp, (0, 1)), (0, 1), (2, 3)) expr = ArrayContraction(ArrayDiagonal(tp, (0, 1)), (0, 1)) assert str(expr) == "ArrayContraction(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))" assert expr.doit() == ArrayDiagonal(ArrayContraction(ArrayTensorProduct(N, M, P), (0, 1)), (0, 1)) expr = ArrayContraction(PermuteDims(M, [1, 0]), (0, 1)) assert str(expr) == "ArrayContraction(PermuteDims(M, (0 1)), (0, 1))" assert expr.doit() == ArrayContraction(M, (0, 1)) # ArrayDiagonal: expr = ArrayDiagonal(ArrayDiagonal(tp, (0, 2)), (0, 1)) assert str(expr) == "ArrayDiagonal(ArrayDiagonal(ArrayTensorProduct(M, N, P), (0, 2)), (0, 1))" assert expr.doit() == ArrayDiagonal(tp, (0, 2), (1, 3)) expr = ArrayDiagonal(ArrayDiagonal(ArrayDiagonal(tp, (0, 1)), (0, 1)), (0, 1)) assert expr.doit() == ArrayDiagonal(tp, (0, 1), (2, 3), (4, 5)) assert expr._canonicalize() == expr.doit() expr = ArrayDiagonal(ArrayContraction(tp, (0, 1)), (0, 1)) assert str(expr) == "ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, P), (0, 1)), (0, 1))" assert expr.doit() == expr expr = ArrayDiagonal(PermuteDims(M, [1, 0]), (0, 1)) assert str(expr) == "ArrayDiagonal(PermuteDims(M, (0 1)), (0, 1))" assert expr.doit() == ArrayDiagonal(M, (0, 1)) # ArrayAdd: expr = ArrayAdd(M) assert isinstance(expr, ArrayAdd) assert expr.doit() == M expr = ArrayAdd(ArrayAdd(M, N), P) assert str(expr) == "ArrayAdd(ArrayAdd(M, N), P)" assert expr.doit() == ArrayAdd(M, N, P) expr = ArrayAdd(M, ArrayAdd(N, ArrayAdd(P, M))) assert expr.doit() == ArrayAdd(M, N, P, M) assert expr._canonicalize() == ArrayAdd(M, N, ArrayAdd(P, M)) expr = ArrayAdd(M, ZeroArray(k, k), N) assert str(expr) == "ArrayAdd(M, ZeroArray(k, k), N)" assert expr.doit() == ArrayAdd(M, N) # PermuteDims: expr = PermuteDims(PermuteDims(M, [1, 0]), [1, 0]) assert str(expr) == "PermuteDims(PermuteDims(M, (0 1)), (0 1))" assert expr.doit() == M expr = PermuteDims(PermuteDims(PermuteDims(M, [1, 0]), [1, 0]), [1, 0]) assert expr.doit() == PermuteDims(M, [1, 0]) assert expr._canonicalize() == expr.doit()
def test_arrayexpr_array_flatten(): # Flatten nested ArrayTensorProduct objects: expr1 = _array_tensor_product(M, N) expr2 = _array_tensor_product(P, Q) expr = _array_tensor_product(expr1, expr2) assert expr == _array_tensor_product(M, N, P, Q) assert expr.args == (M, N, P, Q) # Flatten mixed ArrayTensorProduct and ArrayContraction objects: cg1 = _array_contraction(expr1, (1, 2)) cg2 = _array_contraction(expr2, (0, 3)) expr = _array_tensor_product(cg1, cg2) assert expr == _array_contraction(_array_tensor_product(M, N, P, Q), (1, 2), (4, 7)) expr = _array_tensor_product(M, cg1) assert expr == _array_contraction(_array_tensor_product(M, M, N), (3, 4)) # Flatten nested ArrayContraction objects: cgnested = _array_contraction(cg1, (0, 1)) assert cgnested == _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2)) cgnested = _array_contraction(_array_tensor_product(cg1, cg2), (0, 3)) assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 6), (1, 2), (4, 7)) cg3 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4)) cgnested = _array_contraction(cg3, (0, 1)) assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 5), (1, 3), (2, 4)) cgnested = _array_contraction(cg3, (0, 3), (1, 2)) assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6)) cg4 = _array_contraction(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7)) cgnested = _array_contraction(cg4, (0, 1)) assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7)) cgnested = _array_contraction(cg4, (0, 1), (2, 3)) assert cgnested == _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6)) cg = _array_diagonal(cg4) assert cg == cg4 assert isinstance(cg, type(cg4)) # Flatten nested ArrayDiagonal objects: cg1 = _array_diagonal(expr1, (1, 2)) cg2 = _array_diagonal(expr2, (0, 3)) cg3 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4)) cg4 = _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7)) cgnested = _array_diagonal(cg1, (0, 1)) assert cgnested == _array_diagonal(_array_tensor_product(M, N), (1, 2), (0, 3)) cgnested = _array_diagonal(cg3, (1, 2)) assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 3), (2, 4), (5, 6)) cgnested = _array_diagonal(cg4, (1, 2)) assert cgnested == _array_diagonal(_array_tensor_product(M, N, P, Q), (1, 5), (3, 7), (2, 4)) cg = _array_add(M, N) cg2 = _array_add(cg, P) assert isinstance(cg2, ArrayAdd) assert cg2.args == (M, N, P) assert cg2.shape == (k, k) expr = _array_tensor_product(_array_diagonal(X, (0, 1)), _array_diagonal(A, (0, 1))) assert expr == _array_diagonal(_array_tensor_product(X, A), (0, 1), (2, 3)) expr1 = _array_diagonal(_array_tensor_product(X, A), (1, 2)) expr2 = _array_tensor_product(expr1, a) assert expr2 == _permute_dims(_array_diagonal(_array_tensor_product(X, A, a), (1, 2)), [0, 1, 4, 2, 3]) expr1 = _array_contraction(_array_tensor_product(X, A), (1, 2)) expr2 = _array_tensor_product(expr1, a) assert isinstance(expr2, ArrayContraction) assert isinstance(expr2.expr, ArrayTensorProduct) cg = _array_tensor_product(_array_diagonal(_array_tensor_product(A, X, Y), (0, 3), (1, 5)), a, b) assert cg == _permute_dims(_array_diagonal(_array_tensor_product(A, X, Y, a, b), (0, 3), (1, 5)), [0, 1, 6, 7, 2, 3, 4, 5])