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