I am done

rl/Lib/site-packages/sympy/tensor/array/expressions/__init__.py

@@ -0,0 +1,178 @@
+r"""
+Array expressions are expressions representing N-dimensional arrays, without
+evaluating them. These expressions represent in a certain way abstract syntax
+trees of operations on N-dimensional arrays.
+
+Every N-dimensional array operator has a corresponding array expression object.
+
+Table of correspondences:
+
+=============================== =============================
+         Array operator           Array expression operator
+=============================== =============================
+        tensorproduct                 ArrayTensorProduct
+        tensorcontraction             ArrayContraction
+        tensordiagonal                ArrayDiagonal
+        permutedims                   PermuteDims
+=============================== =============================
+
+Examples
+========
+
+``ArraySymbol`` objects are the N-dimensional equivalent of ``MatrixSymbol``
+objects in the matrix module:
+
+>>> from sympy.tensor.array.expressions import ArraySymbol
+>>> from sympy.abc import i, j, k
+>>> A = ArraySymbol("A", (3, 2, 4))
+>>> A.shape
+(3, 2, 4)
+>>> A[i, j, k]
+A[i, j, k]
+>>> A.as_explicit()
+[[[A[0, 0, 0], A[0, 0, 1], A[0, 0, 2], A[0, 0, 3]],
+  [A[0, 1, 0], A[0, 1, 1], A[0, 1, 2], A[0, 1, 3]]],
+ [[A[1, 0, 0], A[1, 0, 1], A[1, 0, 2], A[1, 0, 3]],
+  [A[1, 1, 0], A[1, 1, 1], A[1, 1, 2], A[1, 1, 3]]],
+ [[A[2, 0, 0], A[2, 0, 1], A[2, 0, 2], A[2, 0, 3]],
+  [A[2, 1, 0], A[2, 1, 1], A[2, 1, 2], A[2, 1, 3]]]]
+
+Component-explicit arrays can be added inside array expressions:
+
+>>> from sympy import Array
+>>> from sympy import tensorproduct
+>>> from sympy.tensor.array.expressions import ArrayTensorProduct
+>>> a = Array([1, 2, 3])
+>>> b = Array([i, j, k])
+>>> expr = ArrayTensorProduct(a, b, b)
+>>> expr
+ArrayTensorProduct([1, 2, 3], [i, j, k], [i, j, k])
+>>> expr.as_explicit() == tensorproduct(a, b, b)
+True
+
+Constructing array expressions from index-explicit forms
+--------------------------------------------------------
+
+Array expressions are index-implicit. This means they do not use any indices to
+represent array operations. The function ``convert_indexed_to_array( ... )``
+may be used to convert index-explicit expressions to array expressions.
+It takes as input two parameters: the index-explicit expression and the order
+of the indices:
+
+>>> from sympy.tensor.array.expressions import convert_indexed_to_array
+>>> from sympy import Sum
+>>> A = ArraySymbol("A", (3, 3))
+>>> B = ArraySymbol("B", (3, 3))
+>>> convert_indexed_to_array(A[i, j], [i, j])
+A
+>>> convert_indexed_to_array(A[i, j], [j, i])
+PermuteDims(A, (0 1))
+>>> convert_indexed_to_array(A[i, j] + B[j, i], [i, j])
+ArrayAdd(A, PermuteDims(B, (0 1)))
+>>> convert_indexed_to_array(Sum(A[i, j]*B[j, k], (j, 0, 2)), [i, k])
+ArrayContraction(ArrayTensorProduct(A, B), (1, 2))
+
+The diagonal of a matrix in the array expression form:
+
+>>> convert_indexed_to_array(A[i, i], [i])
+ArrayDiagonal(A, (0, 1))
+
+The trace of a matrix in the array expression form:
+
+>>> convert_indexed_to_array(Sum(A[i, i], (i, 0, 2)), [i])
+ArrayContraction(A, (0, 1))
+
+Compatibility with matrices
+---------------------------
+
+Array expressions can be mixed with objects from the matrix module:
+
+>>> from sympy import MatrixSymbol
+>>> from sympy.tensor.array.expressions import ArrayContraction
+>>> M = MatrixSymbol("M", 3, 3)
+>>> N = MatrixSymbol("N", 3, 3)
+
+Express the matrix product in the array expression form:
+
+>>> from sympy.tensor.array.expressions import convert_matrix_to_array
+>>> expr = convert_matrix_to_array(M*N)
+>>> expr
+ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
+
+The expression can be converted back to matrix form:
+
+>>> from sympy.tensor.array.expressions import convert_array_to_matrix
+>>> convert_array_to_matrix(expr)
+M*N
+
+Add a second contraction on the remaining axes in order to get the trace of `M \cdot N`:
+
+>>> expr_tr = ArrayContraction(expr, (0, 1))
+>>> expr_tr
+ArrayContraction(ArrayContraction(ArrayTensorProduct(M, N), (1, 2)), (0, 1))
+
+Flatten the expression by calling ``.doit()`` and remove the nested array contraction operations:
+
+>>> expr_tr.doit()
+ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
+
+Get the explicit form of the array expression:
+
+>>> expr.as_explicit()
+[[M[0, 0]*N[0, 0] + M[0, 1]*N[1, 0] + M[0, 2]*N[2, 0], M[0, 0]*N[0, 1] + M[0, 1]*N[1, 1] + M[0, 2]*N[2, 1], M[0, 0]*N[0, 2] + M[0, 1]*N[1, 2] + M[0, 2]*N[2, 2]],
+ [M[1, 0]*N[0, 0] + M[1, 1]*N[1, 0] + M[1, 2]*N[2, 0], M[1, 0]*N[0, 1] + M[1, 1]*N[1, 1] + M[1, 2]*N[2, 1], M[1, 0]*N[0, 2] + M[1, 1]*N[1, 2] + M[1, 2]*N[2, 2]],
+ [M[2, 0]*N[0, 0] + M[2, 1]*N[1, 0] + M[2, 2]*N[2, 0], M[2, 0]*N[0, 1] + M[2, 1]*N[1, 1] + M[2, 2]*N[2, 1], M[2, 0]*N[0, 2] + M[2, 1]*N[1, 2] + M[2, 2]*N[2, 2]]]
+
+Express the trace of a matrix:
+
+>>> from sympy import Trace
+>>> convert_matrix_to_array(Trace(M))
+ArrayContraction(M, (0, 1))
+>>> convert_matrix_to_array(Trace(M*N))
+ArrayContraction(ArrayTensorProduct(M, N), (0, 3), (1, 2))
+
+Express the transposition of a matrix (will be expressed as a permutation of the axes:
+
+>>> convert_matrix_to_array(M.T)
+PermuteDims(M, (0 1))
+
+Compute the derivative array expressions:
+
+>>> from sympy.tensor.array.expressions import array_derive
+>>> d = array_derive(M, M)
+>>> d
+PermuteDims(ArrayTensorProduct(I, I), (3)(1 2))
+
+Verify that the derivative corresponds to the form computed with explicit matrices:
+
+>>> d.as_explicit()
+[[[[1, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 1, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 1], [0, 0, 0], [0, 0, 0]]], [[[0, 0, 0], [1, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 1], [0, 0, 0]]], [[[0, 0, 0], [0, 0, 0], [1, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 1, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 1]]]]
+>>> Me = M.as_explicit()
+>>> Me.diff(Me)
+[[[[1, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 1, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 1], [0, 0, 0], [0, 0, 0]]], [[[0, 0, 0], [1, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 1], [0, 0, 0]]], [[[0, 0, 0], [0, 0, 0], [1, 0, 0]], [[0, 0, 0], [0, 0, 0], [0, 1, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 1]]]]
+
+"""
+
+__all__ = [
+    "ArraySymbol", "ArrayElement", "ZeroArray", "OneArray",
+    "ArrayTensorProduct",
+    "ArrayContraction",
+    "ArrayDiagonal",
+    "PermuteDims",
+    "ArrayAdd",
+    "ArrayElementwiseApplyFunc",
+    "Reshape",
+    "convert_array_to_matrix",
+    "convert_matrix_to_array",
+    "convert_array_to_indexed",
+    "convert_indexed_to_array",
+    "array_derive",
+]
+
+from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayAdd, PermuteDims, ArrayDiagonal, \
+    ArrayContraction, Reshape, ArraySymbol, ArrayElement, ZeroArray, OneArray, ArrayElementwiseApplyFunc
+from sympy.tensor.array.expressions.arrayexpr_derivatives import array_derive
+from sympy.tensor.array.expressions.from_array_to_indexed import convert_array_to_indexed
+from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array
+from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array

rl/Lib/site-packages/sympy/tensor/array/expressions/arrayexpr_derivatives.py

@@ -0,0 +1,194 @@
+import operator
+from functools import reduce, singledispatch
+
+from sympy.core.expr import Expr
+from sympy.core.singleton import S
+from sympy.matrices.expressions.hadamard import HadamardProduct
+from sympy.matrices.expressions.inverse import Inverse
+from sympy.matrices.expressions.matexpr import (MatrixExpr, MatrixSymbol)
+from sympy.matrices.expressions.special import Identity, OneMatrix
+from sympy.matrices.expressions.transpose import Transpose
+from sympy.combinatorics.permutations import _af_invert
+from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction
+from sympy.tensor.array.expressions.array_expressions import (
+    _ArrayExpr, ZeroArray, ArraySymbol, ArrayTensorProduct, ArrayAdd,
+    PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, get_rank,
+    get_shape, ArrayContraction, _array_tensor_product, _array_contraction,
+    _array_diagonal, _array_add, _permute_dims, Reshape)
+from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+
+
+@singledispatch
+def array_derive(expr, x):
+    """
+    Derivatives (gradients) for array expressions.
+    """
+    raise NotImplementedError(f"not implemented for type {type(expr)}")
+
+
+@array_derive.register(Expr)
+def _(expr: Expr, x: _ArrayExpr):
+    return ZeroArray(*x.shape)
+
+
+@array_derive.register(ArrayTensorProduct)
+def _(expr: ArrayTensorProduct, 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 = _array_tensor_product(*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 = list(range(tot1, tot2)) + \
+               list(range(tot1)) + list(range(tot2, tot3)) + \
+               list(range(tot3, tot4))
+        addend = _permute_dims(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 _array_add(*addend_list)
+
+
+@array_derive.register(ArraySymbol)
+def _(expr: ArraySymbol, x: _ArrayExpr):
+    if expr == x:
+        return _permute_dims(
+            ArrayTensorProduct.fromiter(Identity(i) for i in expr.shape),
+            [2*i for i in range(len(expr.shape))] + [2*i+1 for i in range(len(expr.shape))]
+        )
+    return ZeroArray(*(x.shape + expr.shape))
+
+
+@array_derive.register(MatrixSymbol)
+def _(expr: MatrixSymbol, x: _ArrayExpr):
+    m, n = expr.shape
+    if expr == x:
+        return _permute_dims(
+            _array_tensor_product(Identity(m), Identity(n)),
+            [0, 2, 1, 3]
+        )
+    return ZeroArray(*(x.shape + expr.shape))
+
+
+@array_derive.register(Identity)
+def _(expr: Identity, x: _ArrayExpr):
+    return ZeroArray(*(x.shape + expr.shape))
+
+
+@array_derive.register(OneMatrix)
+def _(expr: OneMatrix, x: _ArrayExpr):
+    return ZeroArray(*(x.shape + expr.shape))
+
+
+@array_derive.register(Transpose)
+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 _permute_dims(fd, [0, 1, 3, 2])
+
+
+@array_derive.register(Inverse)
+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
+
+
+@array_derive.register(ElementwiseApplyFunction)
+def _(expr: ElementwiseApplyFunction, x: Expr):
+    assert get_rank(expr) == 2
+    assert get_rank(x) == 2
+    fdiff = expr._get_function_fdiff()
+    dexpr = array_derive(expr.expr, x)
+    tp = _array_tensor_product(
+        ElementwiseApplyFunction(fdiff, expr.expr),
+        dexpr
+    )
+    td = _array_diagonal(
+        tp, (0, 4), (1, 5)
+    )
+    return td
+
+
+@array_derive.register(ArrayElementwiseApplyFunc)
+def _(expr: ArrayElementwiseApplyFunc, x: Expr):
+    fdiff = expr._get_function_fdiff()
+    subexpr = expr.expr
+    dsubexpr = array_derive(subexpr, x)
+    tp = _array_tensor_product(
+        dsubexpr,
+        ArrayElementwiseApplyFunc(fdiff, subexpr)
+    )
+    b = get_rank(x)
+    c = get_rank(expr)
+    diag_indices = [(b + i, b + c + i) for i in range(c)]
+    return _array_diagonal(tp, *diag_indices)
+
+
+@array_derive.register(MatrixExpr)
+def _(expr: MatrixExpr, x: Expr):
+    cg = convert_matrix_to_array(expr)
+    return array_derive(cg, x)
+
+
+@array_derive.register(HadamardProduct)
+def _(expr: HadamardProduct, x: Expr):
+    raise NotImplementedError()
+
+
+@array_derive.register(ArrayContraction)
+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)
+
+
+@array_derive.register(ArrayDiagonal)
+def _(expr: ArrayDiagonal, x: Expr):
+    dsubexpr = array_derive(expr.expr, x)
+    rank_x = len(get_shape(x))
+    diag_indices = [[j + rank_x for j in i] for i in expr.diagonal_indices]
+    return _array_diagonal(dsubexpr, *diag_indices)
+
+
+@array_derive.register(ArrayAdd)
+def _(expr: ArrayAdd, x: Expr):
+    return _array_add(*[array_derive(arg, x) for arg in expr.args])
+
+
+@array_derive.register(PermuteDims)
+def _(expr: PermuteDims, x: Expr):
+    de = array_derive(expr.expr, x)
+    perm = [0, 1] + [i + 2 for i in expr.permutation.array_form]
+    return _permute_dims(de, perm)
+
+
+@array_derive.register(Reshape)
+def _(expr: Reshape, x: Expr):
+    de = array_derive(expr.expr, x)
+    return Reshape(de, get_shape(x) + expr.shape)
+
+
+def matrix_derive(expr, x):
+    from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+    ce = convert_matrix_to_array(expr)
+    dce = array_derive(ce, x)
+    return convert_array_to_matrix(dce).doit()

rl/Lib/site-packages/sympy/tensor/array/expressions/conv_array_to_indexed.py

@@ -0,0 +1,12 @@
+from sympy.tensor.array.expressions import from_array_to_indexed
+from sympy.utilities.decorator import deprecated
+
+
+_conv_to_from_decorator = deprecated(
+    "module has been renamed by replacing 'conv_' with 'from_' in its name",
+    deprecated_since_version="1.11",
+    active_deprecations_target="deprecated-conv-array-expr-module-names",
+)
+
+
+convert_array_to_indexed = _conv_to_from_decorator(from_array_to_indexed.convert_array_to_indexed)

rl/Lib/site-packages/sympy/tensor/array/expressions/conv_array_to_matrix.py

@@ -0,0 +1,6 @@
+from sympy.tensor.array.expressions import from_array_to_matrix
+from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
+
+convert_array_to_matrix = _conv_to_from_decorator(from_array_to_matrix.convert_array_to_matrix)
+_array2matrix = _conv_to_from_decorator(from_array_to_matrix._array2matrix)
+_remove_trivial_dims = _conv_to_from_decorator(from_array_to_matrix._remove_trivial_dims)

rl/Lib/site-packages/sympy/tensor/array/expressions/conv_indexed_to_array.py

@@ -0,0 +1,4 @@
+from sympy.tensor.array.expressions import from_indexed_to_array
+from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
+
+convert_indexed_to_array = _conv_to_from_decorator(from_indexed_to_array.convert_indexed_to_array)

rl/Lib/site-packages/sympy/tensor/array/expressions/conv_matrix_to_array.py

@@ -0,0 +1,4 @@
+from sympy.tensor.array.expressions import from_matrix_to_array
+from sympy.tensor.array.expressions.conv_array_to_indexed import _conv_to_from_decorator
+
+convert_matrix_to_array = _conv_to_from_decorator(from_matrix_to_array.convert_matrix_to_array)

rl/Lib/site-packages/sympy/tensor/array/expressions/from_array_to_indexed.py

@@ -0,0 +1,84 @@
+import collections.abc
+import operator
+from itertools import accumulate
+
+from sympy import Mul, Sum, Dummy, Add
+from sympy.tensor.array.expressions import PermuteDims, ArrayAdd, ArrayElementwiseApplyFunc, Reshape
+from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, get_rank, ArrayContraction, \
+    ArrayDiagonal, get_shape, _get_array_element_or_slice, _ArrayExpr
+from sympy.tensor.array.expressions.utils import _apply_permutation_to_list
+
+
+def convert_array_to_indexed(expr, indices):
+    return _ConvertArrayToIndexed().do_convert(expr, indices)
+
+
+class _ConvertArrayToIndexed:
+
+    def __init__(self):
+        self.count_dummies = 0
+
+    def do_convert(self, expr, indices):
+        if isinstance(expr, ArrayTensorProduct):
+            cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args]))
+            indices_grp = [indices[cumul[i]:cumul[i+1]] for i in range(len(expr.args))]
+            return Mul.fromiter(self.do_convert(arg, ind) for arg, ind in zip(expr.args, indices_grp))
+        if isinstance(expr, ArrayContraction):
+            new_indices = [None for i in range(get_rank(expr.expr))]
+            limits = []
+            bottom_shape = get_shape(expr.expr)
+            for contraction_index_grp in expr.contraction_indices:
+                d = Dummy(f"d{self.count_dummies}")
+                self.count_dummies += 1
+                dim = bottom_shape[contraction_index_grp[0]]
+                limits.append((d, 0, dim-1))
+                for i in contraction_index_grp:
+                    new_indices[i] = d
+            j = 0
+            for i in range(len(new_indices)):
+                if new_indices[i] is None:
+                    new_indices[i] = indices[j]
+                    j += 1
+            newexpr = self.do_convert(expr.expr, new_indices)
+            return Sum(newexpr, *limits)
+        if isinstance(expr, ArrayDiagonal):
+            new_indices = [None for i in range(get_rank(expr.expr))]
+            ind_pos = expr._push_indices_down(expr.diagonal_indices, list(range(len(indices))), get_rank(expr))
+            for i, index in zip(ind_pos, indices):
+                if isinstance(i, collections.abc.Iterable):
+                    for j in i:
+                        new_indices[j] = index
+                else:
+                    new_indices[i] = index
+            newexpr = self.do_convert(expr.expr, new_indices)
+            return newexpr
+        if isinstance(expr, PermuteDims):
+            permuted_indices = _apply_permutation_to_list(expr.permutation, indices)
+            return self.do_convert(expr.expr, permuted_indices)
+        if isinstance(expr, ArrayAdd):
+            return Add.fromiter(self.do_convert(arg, indices) for arg in expr.args)
+        if isinstance(expr, _ArrayExpr):
+            return expr.__getitem__(tuple(indices))
+        if isinstance(expr, ArrayElementwiseApplyFunc):
+            return expr.function(self.do_convert(expr.expr, indices))
+        if isinstance(expr, Reshape):
+            shape_up = expr.shape
+            shape_down = get_shape(expr.expr)
+            cumul = list(accumulate([1] + list(reversed(shape_up)), operator.mul))
+            one_index = Add.fromiter(i*s for i, s in zip(reversed(indices), cumul))
+            dest_indices = [None for _ in shape_down]
+            c = 1
+            for i, e in enumerate(reversed(shape_down)):
+                if c == 1:
+                    if i == len(shape_down) - 1:
+                        dest_indices[i] = one_index
+                    else:
+                        dest_indices[i] = one_index % e
+                elif i == len(shape_down) - 1:
+                    dest_indices[i] = one_index // c
+                else:
+                    dest_indices[i] = one_index // c % e
+                c *= e
+            dest_indices.reverse()
+            return self.do_convert(expr.expr, dest_indices)
+        return _get_array_element_or_slice(expr, indices)

rl/Lib/site-packages/sympy/tensor/array/expressions/from_indexed_to_array.py

@@ -0,0 +1,257 @@
+from collections import defaultdict
+
+from sympy import Function
+from sympy.combinatorics.permutations import _af_invert
+from sympy.concrete.summations import Sum
+from sympy.core.add import Add
+from sympy.core.mul import Mul
+from sympy.core.numbers import Integer
+from sympy.core.power import Pow
+from sympy.core.sorting import default_sort_key
+from sympy.functions.special.tensor_functions import KroneckerDelta
+from sympy.tensor.array.expressions import ArrayElementwiseApplyFunc
+from sympy.tensor.indexed import (Indexed, IndexedBase)
+from sympy.combinatorics import Permutation
+from sympy.matrices.expressions.matexpr import MatrixElement
+from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, \
+    get_shape, ArrayElement, _array_tensor_product, _array_diagonal, _array_contraction, _array_add, \
+    _permute_dims, OneArray, ArrayAdd
+from sympy.tensor.array.expressions.utils import _get_argindex, _get_diagonal_indices
+
+
+def convert_indexed_to_array(expr, first_indices=None):
+    r"""
+    Parse indexed expression into a form useful for code generation.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array
+    >>> from sympy import MatrixSymbol, Sum, symbols
+
+    >>> i, j, k, d = symbols("i j k d")
+    >>> M = MatrixSymbol("M", d, d)
+    >>> N = MatrixSymbol("N", d, d)
+
+    Recognize the trace in summation form:
+
+    >>> expr = Sum(M[i, i], (i, 0, d-1))
+    >>> convert_indexed_to_array(expr)
+    ArrayContraction(M, (0, 1))
+
+    Recognize the extraction of the diagonal by using the same index `i` on
+    both axes of the matrix:
+
+    >>> expr = M[i, i]
+    >>> convert_indexed_to_array(expr)
+    ArrayDiagonal(M, (0, 1))
+
+    This function can help perform the transformation expressed in two
+    different mathematical notations as:
+
+    `\sum_{j=0}^{N-1} A_{i,j} B_{j,k} \Longrightarrow \mathbf{A}\cdot \mathbf{B}`
+
+    Recognize the matrix multiplication in summation form:
+
+    >>> expr = Sum(M[i, j]*N[j, k], (j, 0, d-1))
+    >>> convert_indexed_to_array(expr)
+    ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
+
+    Specify that ``k`` has to be the starting index:
+
+    >>> convert_indexed_to_array(expr, first_indices=[k])
+    ArrayContraction(ArrayTensorProduct(N, M), (0, 3))
+    """
+
+    result, indices = _convert_indexed_to_array(expr)
+
+    if any(isinstance(i, (int, Integer)) for i in indices):
+        result = ArrayElement(result, indices)
+        indices = []
+
+    if not first_indices:
+        return result
+
+    def _check_is_in(elem, indices):
+        if elem in indices:
+            return True
+        if any(elem in i for i in indices if isinstance(i, frozenset)):
+            return True
+        return False
+
+    repl = {j: i for i in indices if isinstance(i, frozenset) for j in i}
+    first_indices = [repl.get(i, i) for i in first_indices]
+    for i in first_indices:
+        if not _check_is_in(i, indices):
+            first_indices.remove(i)
+    first_indices.extend([i for i in indices if not _check_is_in(i, first_indices)])
+
+    def _get_pos(elem, indices):
+        if elem in indices:
+            return indices.index(elem)
+        for i, e in enumerate(indices):
+            if not isinstance(e, frozenset):
+                continue
+            if elem in e:
+                return i
+        raise ValueError("not found")
+
+    permutation = _af_invert([_get_pos(i, first_indices) for i in indices])
+    if isinstance(result, ArrayAdd):
+        return _array_add(*[_permute_dims(arg, permutation) for arg in result.args])
+    else:
+        return _permute_dims(result, permutation)
+
+
+def _convert_indexed_to_array(expr):
+    if isinstance(expr, Sum):
+        function = expr.function
+        summation_indices = expr.variables
+        subexpr, subindices = _convert_indexed_to_array(function)
+        subindicessets = {j: i for i in subindices if isinstance(i, frozenset) for j in i}
+        summation_indices = sorted({subindicessets.get(i, i) for i in summation_indices}, key=default_sort_key)
+        # TODO: check that Kronecker delta is only contracted to one other element:
+        kronecker_indices = set()
+        if isinstance(function, Mul):
+            for arg in function.args:
+                if not isinstance(arg, KroneckerDelta):
+                    continue
+                arg_indices = sorted(set(arg.indices), key=default_sort_key)
+                if len(arg_indices) == 2:
+                    kronecker_indices.update(arg_indices)
+        kronecker_indices = sorted(kronecker_indices, key=default_sort_key)
+        # Check dimensional consistency:
+        shape = get_shape(subexpr)
+        if shape:
+            for ind, istart, iend in expr.limits:
+                i = _get_argindex(subindices, ind)
+                if istart != 0 or iend+1 != shape[i]:
+                    raise ValueError("summation index and array dimension mismatch: %s" % ind)
+        contraction_indices = []
+        subindices = list(subindices)
+        if isinstance(subexpr, ArrayDiagonal):
+            diagonal_indices = list(subexpr.diagonal_indices)
+            dindices = subindices[-len(diagonal_indices):]
+            subindices = subindices[:-len(diagonal_indices)]
+            for index in summation_indices:
+                if index in dindices:
+                    position = dindices.index(index)
+                    contraction_indices.append(diagonal_indices[position])
+                    diagonal_indices[position] = None
+            diagonal_indices = [i for i in diagonal_indices if i is not None]
+            for i, ind in enumerate(subindices):
+                if ind in summation_indices:
+                    pass
+            if diagonal_indices:
+                subexpr = _array_diagonal(subexpr.expr, *diagonal_indices)
+            else:
+                subexpr = subexpr.expr
+
+        axes_contraction = defaultdict(list)
+        for i, ind in enumerate(subindices):
+            include = all(j not in kronecker_indices for j in ind) if isinstance(ind, frozenset) else ind not in kronecker_indices
+            if ind in summation_indices and include:
+                axes_contraction[ind].append(i)
+                subindices[i] = None
+        for k, v in axes_contraction.items():
+            if any(i in kronecker_indices for i in k) if isinstance(k, frozenset) else k in kronecker_indices:
+                continue
+            contraction_indices.append(tuple(v))
+        free_indices = [i for i in subindices if i is not None]
+        indices_ret = list(free_indices)
+        indices_ret.sort(key=lambda x: free_indices.index(x))
+        return _array_contraction(
+                subexpr,
+                *contraction_indices,
+                free_indices=free_indices
+            ), tuple(indices_ret)
+    if isinstance(expr, Mul):
+        args, indices = zip(*[_convert_indexed_to_array(arg) for arg in expr.args])
+        # Check if there are KroneckerDelta objects:
+        kronecker_delta_repl = {}
+        for arg in args:
+            if not isinstance(arg, KroneckerDelta):
+                continue
+            # Diagonalize two indices:
+            i, j = arg.indices
+            kindices = set(arg.indices)
+            if i in kronecker_delta_repl:
+                kindices.update(kronecker_delta_repl[i])
+            if j in kronecker_delta_repl:
+                kindices.update(kronecker_delta_repl[j])
+            kindices = frozenset(kindices)
+            for index in kindices:
+                kronecker_delta_repl[index] = kindices
+        # Remove KroneckerDelta objects, their relations should be handled by
+        # ArrayDiagonal:
+        newargs = []
+        newindices = []
+        for arg, loc_indices in zip(args, indices):
+            if isinstance(arg, KroneckerDelta):
+                continue
+            newargs.append(arg)
+            newindices.append(loc_indices)
+        flattened_indices = [kronecker_delta_repl.get(j, j) for i in newindices for j in i]
+        diagonal_indices, ret_indices = _get_diagonal_indices(flattened_indices)
+        tp = _array_tensor_product(*newargs)
+        if diagonal_indices:
+            return _array_diagonal(tp, *diagonal_indices), ret_indices
+        else:
+            return tp, ret_indices
+    if isinstance(expr, MatrixElement):
+        indices = expr.args[1:]
+        diagonal_indices, ret_indices = _get_diagonal_indices(indices)
+        if diagonal_indices:
+            return _array_diagonal(expr.args[0], *diagonal_indices), ret_indices
+        else:
+            return expr.args[0], ret_indices
+    if isinstance(expr, ArrayElement):
+        indices = expr.indices
+        diagonal_indices, ret_indices = _get_diagonal_indices(indices)
+        if diagonal_indices:
+            return _array_diagonal(expr.name, *diagonal_indices), ret_indices
+        else:
+            return expr.name, ret_indices
+    if isinstance(expr, Indexed):
+        indices = expr.indices
+        diagonal_indices, ret_indices = _get_diagonal_indices(indices)
+        if diagonal_indices:
+            return _array_diagonal(expr.base, *diagonal_indices), ret_indices
+        else:
+            return expr.args[0], ret_indices
+    if isinstance(expr, IndexedBase):
+        raise NotImplementedError
+    if isinstance(expr, KroneckerDelta):
+        return expr, expr.indices
+    if isinstance(expr, Add):
+        args, indices = zip(*[_convert_indexed_to_array(arg) for arg in expr.args])
+        args = list(args)
+        # Check if all indices are compatible. Otherwise expand the dimensions:
+        index0 = []
+        shape0 = []
+        for arg, arg_indices in zip(args, indices):
+            arg_indices_set = set(arg_indices)
+            arg_indices_missing = arg_indices_set.difference(index0)
+            index0.extend([i for i in arg_indices if i in arg_indices_missing])
+            arg_shape = get_shape(arg)
+            shape0.extend([arg_shape[i] for i, e in enumerate(arg_indices) if e in arg_indices_missing])
+        for i, (arg, arg_indices) in enumerate(zip(args, indices)):
+            if len(arg_indices) < len(index0):
+                missing_indices_pos = [i for i, e in enumerate(index0) if e not in arg_indices]
+                missing_shape = [shape0[i] for i in missing_indices_pos]
+                arg_indices = tuple(index0[j] for j in missing_indices_pos) + arg_indices
+                args[i] = _array_tensor_product(OneArray(*missing_shape), args[i])
+            permutation = Permutation([arg_indices.index(j) for j in index0])
+            # Perform index permutations:
+            args[i] = _permute_dims(args[i], permutation)
+        return _array_add(*args), tuple(index0)
+    if isinstance(expr, Pow):
+        subexpr, subindices = _convert_indexed_to_array(expr.base)
+        if isinstance(expr.exp, (int, Integer)):
+            diags = zip(*[(2*i, 2*i + 1) for i in range(expr.exp)])
+            arr = _array_diagonal(_array_tensor_product(*[subexpr for i in range(expr.exp)]), *diags)
+            return arr, subindices
+    if isinstance(expr, Function):
+        subexpr, subindices = _convert_indexed_to_array(expr.args[0])
+        return ArrayElementwiseApplyFunc(type(expr), subexpr), subindices
+    return expr, ()

rl/Lib/site-packages/sympy/tensor/array/expressions/from_matrix_to_array.py

@@ -0,0 +1,87 @@
+from sympy import KroneckerProduct
+from sympy.core.basic import Basic
+from sympy.core.function import Lambda
+from sympy.core.mul import Mul
+from sympy.core.numbers import Integer
+from sympy.core.power import Pow
+from sympy.core.singleton import S
+from sympy.core.symbol import (Dummy, symbols)
+from sympy.matrices.expressions.hadamard import (HadamardPower, HadamardProduct)
+from sympy.matrices.expressions.matadd import MatAdd
+from sympy.matrices.expressions.matmul import MatMul
+from sympy.matrices.expressions.matpow import MatPow
+from sympy.matrices.expressions.trace import Trace
+from sympy.matrices.expressions.transpose import Transpose
+from sympy.matrices.expressions.matexpr import MatrixExpr
+from sympy.tensor.array.expressions.array_expressions import \
+    ArrayElementwiseApplyFunc, _array_tensor_product, _array_contraction, \
+    _array_diagonal, _array_add, _permute_dims, Reshape
+
+
+def convert_matrix_to_array(expr: Basic) -> Basic:
+    if isinstance(expr, MatMul):
+        args_nonmat = []
+        args = []
+        for arg in expr.args:
+            if isinstance(arg, MatrixExpr):
+                args.append(arg)
+            else:
+                args_nonmat.append(convert_matrix_to_array(arg))
+        contractions = [(2*i+1, 2*i+2) for i in range(len(args)-1)]
+        scalar = _array_tensor_product(*args_nonmat) if args_nonmat else S.One
+        if scalar == 1:
+            tprod = _array_tensor_product(
+                *[convert_matrix_to_array(arg) for arg in args])
+        else:
+            tprod = _array_tensor_product(
+                scalar,
+                *[convert_matrix_to_array(arg) for arg in args])
+        return _array_contraction(
+                tprod,
+                *contractions
+        )
+    elif isinstance(expr, MatAdd):
+        return _array_add(
+                *[convert_matrix_to_array(arg) for arg in expr.args]
+        )
+    elif isinstance(expr, Transpose):
+        return _permute_dims(
+                convert_matrix_to_array(expr.args[0]), [1, 0]
+        )
+    elif isinstance(expr, Trace):
+        inner_expr: MatrixExpr = convert_matrix_to_array(expr.arg) # type: ignore
+        return _array_contraction(inner_expr, (0, len(inner_expr.shape) - 1))
+    elif isinstance(expr, Mul):
+        return _array_tensor_product(*[convert_matrix_to_array(i) for i in expr.args])
+    elif isinstance(expr, Pow):
+        base = convert_matrix_to_array(expr.base)
+        if (expr.exp > 0) == True:
+            return _array_tensor_product(*[base for i in range(expr.exp)])
+        else:
+            return expr
+    elif isinstance(expr, MatPow):
+        base = convert_matrix_to_array(expr.base)
+        if expr.exp.is_Integer != True:
+            b = symbols("b", cls=Dummy)
+            return ArrayElementwiseApplyFunc(Lambda(b, b**expr.exp), convert_matrix_to_array(base))
+        elif (expr.exp > 0) == True:
+            return convert_matrix_to_array(MatMul.fromiter(base for i in range(expr.exp)))
+        else:
+            return expr
+    elif isinstance(expr, HadamardProduct):
+        tp = _array_tensor_product(*[convert_matrix_to_array(arg) for arg in expr.args])
+        diag = [[2*i for i in range(len(expr.args))], [2*i+1 for i in range(len(expr.args))]]
+        return _array_diagonal(tp, *diag)
+    elif isinstance(expr, HadamardPower):
+        base, exp = expr.args
+        if isinstance(exp, Integer) and exp > 0:
+            return convert_matrix_to_array(HadamardProduct.fromiter(base for i in range(exp)))
+        else:
+            d = Dummy("d")
+            return ArrayElementwiseApplyFunc(Lambda(d, d**exp), base)
+    elif isinstance(expr, KroneckerProduct):
+        kp_args = [convert_matrix_to_array(arg) for arg in expr.args]
+        permutation = [2*i for i in range(len(kp_args))] + [2*i + 1 for i in range(len(kp_args))]
+        return Reshape(_permute_dims(_array_tensor_product(*kp_args), permutation), expr.shape)
+    else:
+        return expr

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_array_expressions.py

@@ -0,0 +1,808 @@
+import random
+
+from sympy import tensordiagonal, eye, KroneckerDelta, Array
+from sympy.core.symbol import symbols
+from sympy.functions.elementary.trigonometric import (cos, sin)
+from sympy.matrices.expressions.diagonal import DiagMatrix
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.matrices.expressions.special import ZeroMatrix
+from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensorproduct)
+from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
+from sympy.combinatorics import Permutation
+from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, ArraySymbol, ArrayElement, \
+    PermuteDims, ArrayContraction, ArrayTensorProduct, ArrayDiagonal, \
+    ArrayAdd, nest_permutation, ArrayElementwiseApplyFunc, _EditArrayContraction, _ArgE, _array_tensor_product, \
+    _array_contraction, _array_diagonal, _array_add, _permute_dims, Reshape
+from sympy.testing.pytest import raises
+
+i, j, k, l, m, n = symbols("i j k l m n")
+
+
+M = ArraySymbol("M", (k, k))
+N = ArraySymbol("N", (k, k))
+P = ArraySymbol("P", (k, k))
+Q = ArraySymbol("Q", (k, k))
+
+A = ArraySymbol("A", (k, k))
+B = ArraySymbol("B", (k, k))
+C = ArraySymbol("C", (k, k))
+D = ArraySymbol("D", (k, k))
+
+X = ArraySymbol("X", (k, k))
+Y = ArraySymbol("Y", (k, k))
+
+a = ArraySymbol("a", (k, 1))
+b = ArraySymbol("b", (k, 1))
+c = ArraySymbol("c", (k, 1))
+d = ArraySymbol("d", (k, 1))
+
+
+def test_array_symbol_and_element():
+    A = ArraySymbol("A", (2,))
+    A0 = ArrayElement(A, (0,))
+    A1 = ArrayElement(A, (1,))
+    assert A[0] == A0
+    assert A[1] != A0
+    assert A.as_explicit() == ImmutableDenseNDimArray([A0, A1])
+
+    A2 = tensorproduct(A, A)
+    assert A2.shape == (2, 2)
+    # TODO: not yet supported:
+    # assert A2.as_explicit() == Array([[A[0]*A[0], A[1]*A[0]], [A[0]*A[1], A[1]*A[1]]])
+    A3 = tensorcontraction(A2, (0, 1))
+    assert A3.shape == ()
+    # TODO: not yet supported:
+    # assert A3.as_explicit() == Array([])
+
+    A = ArraySymbol("A", (2, 3, 4))
+    Ae = A.as_explicit()
+    assert Ae == ImmutableDenseNDimArray(
+        [[[ArrayElement(A, (i, j, k)) for k in range(4)] for j in range(3)] for i in range(2)])
+
+    p = _permute_dims(A, Permutation(0, 2, 1))
+    assert isinstance(p, PermuteDims)
+
+    A = ArraySymbol("A", (2,))
+    raises(IndexError, lambda: A[()])
+    raises(IndexError, lambda: A[0, 1])
+    raises(ValueError, lambda: A[-1])
+    raises(ValueError, lambda: A[2])
+
+    O = OneArray(3, 4)
+    Z = ZeroArray(m, n)
+
+    raises(IndexError, lambda: O[()])
+    raises(IndexError, lambda: O[1, 2, 3])
+    raises(ValueError, lambda: O[3, 0])
+    raises(ValueError, lambda: O[0, 4])
+
+    assert O[1, 2] == 1
+    assert Z[1, 2] == 0
+
+
+def test_zero_array():
+    assert ZeroArray() == 0
+    assert ZeroArray().is_Integer
+
+    za = ZeroArray(3, 2, 4)
+    assert za.shape == (3, 2, 4)
+    za_e = za.as_explicit()
+    assert za_e.shape == (3, 2, 4)
+
+    m, n, k = symbols("m n k")
+    za = ZeroArray(m, n, k, 2)
+    assert za.shape == (m, n, k, 2)
+    raises(ValueError, lambda: za.as_explicit())
+
+
+def test_one_array():
+    assert OneArray() == 1
+    assert OneArray().is_Integer
+
+    oa = OneArray(3, 2, 4)
+    assert oa.shape == (3, 2, 4)
+    oa_e = oa.as_explicit()
+    assert oa_e.shape == (3, 2, 4)
+
+    m, n, k = symbols("m n k")
+    oa = OneArray(m, n, k, 2)
+    assert oa.shape == (m, n, k, 2)
+    raises(ValueError, lambda: oa.as_explicit())
+
+
+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_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])
+
+
+def test_arrayexpr_array_diagonal():
+    cg = _array_diagonal(M, (1, 0))
+    assert cg == _array_diagonal(M, (0, 1))
+
+    cg = _array_diagonal(_array_tensor_product(M, N, P), (4, 1), (2, 0))
+    assert cg == _array_diagonal(_array_tensor_product(M, N, P), (1, 4), (0, 2))
+
+    cg = _array_diagonal(_array_tensor_product(M, N), (1, 2), (3,), allow_trivial_diags=True)
+    assert cg == _permute_dims(_array_diagonal(_array_tensor_product(M, N), (1, 2)), [0, 2, 1])
+
+    Ax = ArraySymbol("Ax", shape=(1, 2, 3, 4, 3, 5, 6, 2, 7))
+    cg = _array_diagonal(Ax, (1, 7), (3,), (2, 4), (6,), allow_trivial_diags=True)
+    assert cg == _permute_dims(_array_diagonal(Ax, (1, 7), (2, 4)), [0, 2, 4, 5, 1, 6, 3])
+
+    cg = _array_diagonal(M, (0,), allow_trivial_diags=True)
+    assert cg == _permute_dims(M, [1, 0])
+
+    raises(ValueError, lambda: _array_diagonal(M, (0, 0)))
+
+
+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_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_push_indices_up_and_down():
+
+    indices = list(range(12))
+
+    contr_diag_indices = [(0, 6), (2, 8)]
+    assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (1, 3, 4, 5, 7, 9, 10, 11, 12, 13, 14, 15)
+    assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (None, 0, None, 1, 2, 3, None, 4, None, 5, 6, 7)
+
+    assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (1, 3, 4, 5, 7, 9, (0, 6), (2, 8), None, None, None, None)
+    assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (6, 0, 7, 1, 2, 3, 6, 4, 7, 5, None, None)
+
+    contr_diag_indices = [(1, 2), (7, 8)]
+    assert ArrayContraction._push_indices_down(contr_diag_indices, indices) == (0, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15)
+    assert ArrayContraction._push_indices_up(contr_diag_indices, indices) == (0, None, None, 1, 2, 3, 4, None, None, 5, 6, 7)
+
+    assert ArrayDiagonal._push_indices_down(contr_diag_indices, indices, 10) == (0, 3, 4, 5, 6, 9, (1, 2), (7, 8), None, None, None, None)
+    assert ArrayDiagonal._push_indices_up(contr_diag_indices, indices, 10) == (0, 6, 6, 1, 2, 3, 4, 7, 7, 5, None, None)
+
+
+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_nested_permutations():
+
+    cg = _permute_dims(_permute_dims(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 = _permute_dims(
+            _permute_dims(
+                _array_tensor_product(M, N, P),
+                p1),
+            p2
+        )
+        result = _permute_dims(
+            _array_tensor_product(M, N, P),
+            p2*p1
+        )
+        assert cg == result
+
+        # Check that `permutedims` behaves the same way with explicit-component arrays:
+        result1 = _permute_dims(_permute_dims(cge, p1), p2)
+        result2 = _permute_dims(cge, p2*p1)
+        assert result1 == result2
+
+
+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_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_canonicalize_diagonal__permute_dims():
+    tp = _array_tensor_product(M, Q, N, P)
+    expr = _array_diagonal(
+        _permute_dims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7),
+        (0, 3))
+    result = _array_diagonal(tp, (2, 6, 7), (3, 5), (0, 4))
+    assert expr == result
+
+    tp = _array_tensor_product(M, N, P, Q)
+    expr = _array_diagonal(_permute_dims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4))
+    result = _array_diagonal(_array_tensor_product(M, P, N, Q), (3, 4, 5), (1, 2))
+    assert expr == result
+
+
+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_array_wrong_permutation_size():
+    cg = _array_tensor_product(M, N)
+    raises(ValueError, lambda: _permute_dims(cg, [1, 0]))
+    raises(ValueError, lambda: _permute_dims(cg, [1, 0, 2, 3, 5, 4]))
+
+
+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 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_array_expr_applyfunc():
+
+    A = ArraySymbol("A", (3, k, 2))
+    aaf = ArrayElementwiseApplyFunc(sin, A)
+    assert aaf.shape == (3, k, 2)
+
+
+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_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()
+
+    # Reshape
+
+    expr = Reshape(A, (k**2,))
+    assert expr.shape == (k**2,)
+    assert isinstance(expr, Reshape)
+
+
+def test_array_expr_construction_with_functions():
+
+    tp = tensorproduct(M, N)
+    assert tp == ArrayTensorProduct(M, N)
+
+    expr = tensorproduct(A, eye(2))
+    assert expr == ArrayTensorProduct(A, eye(2))
+
+    # Contraction:
+
+    expr = tensorcontraction(M, (0, 1))
+    assert expr == ArrayContraction(M, (0, 1))
+
+    expr = tensorcontraction(tp, (1, 2))
+    assert expr == ArrayContraction(tp, (1, 2))
+
+    expr = tensorcontraction(tensorcontraction(tp, (1, 2)), (0, 1))
+    assert expr == ArrayContraction(tp, (0, 3), (1, 2))
+
+    # Diagonalization:
+
+    expr = tensordiagonal(M, (0, 1))
+    assert expr == ArrayDiagonal(M, (0, 1))
+
+    expr = tensordiagonal(tensordiagonal(tp, (0, 1)), (0, 1))
+    assert expr == ArrayDiagonal(tp, (0, 1), (2, 3))
+
+    # Permutation of dimensions:
+
+    expr = permutedims(M, [1, 0])
+    assert expr == PermuteDims(M, [1, 0])
+
+    expr = permutedims(PermuteDims(tp, [1, 0, 2, 3]), [0, 1, 3, 2])
+    assert expr == PermuteDims(tp, [1, 0, 3, 2])
+
+    expr = PermuteDims(tp, index_order_new=["a", "b", "c", "d"], index_order_old=["d", "c", "b", "a"])
+    assert expr == PermuteDims(tp, [3, 2, 1, 0])
+
+    arr = Array(range(32)).reshape(2, 2, 2, 2, 2)
+    expr = PermuteDims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
+    assert expr == PermuteDims(arr, [2, 0, 4, 3, 1])
+    assert expr.as_explicit() == permutedims(arr, index_order_new=["a", "b", "c", "d", "e"], index_order_old=['b', 'e', 'a', 'd', 'c'])
+
+
+def test_array_element_expressions():
+    # Check commutative property:
+    assert M[0, 0]*N[0, 0] == N[0, 0]*M[0, 0]
+
+    # Check derivatives:
+    assert M[0, 0].diff(M[0, 0]) == 1
+    assert M[0, 0].diff(M[1, 0]) == 0
+    assert M[0, 0].diff(N[0, 0]) == 0
+    assert M[0, 1].diff(M[i, j]) == KroneckerDelta(i, 0)*KroneckerDelta(j, 1)
+    assert M[0, 1].diff(N[i, j]) == 0
+
+    K4 = ArraySymbol("K4", shape=(k, k, k, k))
+
+    assert K4[i, j, k, l].diff(K4[1, 2, 3, 4]) == (
+        KroneckerDelta(i, 1)*KroneckerDelta(j, 2)*KroneckerDelta(k, 3)*KroneckerDelta(l, 4)
+    )
+
+
+def test_array_expr_reshape():
+
+    A = MatrixSymbol("A", 2, 2)
+    B = ArraySymbol("B", (2, 2, 2))
+    C = Array([1, 2, 3, 4])
+
+    expr = Reshape(A, (4,))
+    assert expr.expr == A
+    assert expr.shape == (4,)
+    assert expr.as_explicit() == Array([A[0, 0], A[0, 1], A[1, 0], A[1, 1]])
+
+    expr = Reshape(B, (2, 4))
+    assert expr.expr == B
+    assert expr.shape == (2, 4)
+    ee = expr.as_explicit()
+    assert isinstance(ee, ImmutableDenseNDimArray)
+    assert ee.shape == (2, 4)
+    assert ee == Array([[B[0, 0, 0], B[0, 0, 1], B[0, 1, 0], B[0, 1, 1]], [B[1, 0, 0], B[1, 0, 1], B[1, 1, 0], B[1, 1, 1]]])
+
+    expr = Reshape(A, (k, 2))
+    assert expr.shape == (k, 2)
+
+    raises(ValueError, lambda: Reshape(A, (2, 3)))
+    raises(ValueError, lambda: Reshape(A, (3,)))
+
+    expr = Reshape(C, (2, 2))
+    assert expr.expr == C
+    assert expr.shape == (2, 2)
+    assert expr.doit() == Array([[1, 2], [3, 4]])
+
+
+def test_array_expr_as_explicit_with_explicit_component_arrays():
+    # Test if .as_explicit() works with explicit-component arrays
+    # nested in array expressions:
+    from sympy.abc import x, y, z, t
+    A = Array([[x, y], [z, t]])
+    assert ArrayTensorProduct(A, A).as_explicit() == tensorproduct(A, A)
+    assert ArrayDiagonal(A, (0, 1)).as_explicit() == tensordiagonal(A, (0, 1))
+    assert ArrayContraction(A, (0, 1)).as_explicit() == tensorcontraction(A, (0, 1))
+    assert ArrayAdd(A, A).as_explicit() == A + A
+    assert ArrayElementwiseApplyFunc(sin, A).as_explicit() == A.applyfunc(sin)
+    assert PermuteDims(A, [1, 0]).as_explicit() == permutedims(A, [1, 0])
+    assert Reshape(A, [4]).as_explicit() == A.reshape(4)

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_arrayexpr_derivatives.py

@@ -0,0 +1,78 @@
+from sympy.core.symbol import symbols
+from sympy.functions.elementary.trigonometric import (cos, sin)
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.matrices.expressions.special import Identity
+from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction
+from sympy.tensor.array.expressions.array_expressions import ArraySymbol, ArrayTensorProduct, \
+    PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, ArrayContraction, _permute_dims, Reshape
+from sympy.tensor.array.expressions.arrayexpr_derivatives import array_derive
+
+k = symbols("k")
+
+I = Identity(k)
+X = MatrixSymbol("X", k, k)
+x = MatrixSymbol("x", k, 1)
+
+A = MatrixSymbol("A", k, k)
+B = MatrixSymbol("B", k, k)
+C = MatrixSymbol("C", k, k)
+D = MatrixSymbol("D", k, k)
+
+A1 = ArraySymbol("A", (3, 2, k))
+
+
+def test_arrayexpr_derivatives1():
+
+    res = array_derive(X, X)
+    assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
+
+    cg = ArrayTensorProduct(A, X, B)
+    res = array_derive(cg, X)
+    assert res == _permute_dims(
+        ArrayTensorProduct(I, A, I, B),
+        [0, 4, 2, 3, 1, 5, 6, 7])
+
+    cg = ArrayContraction(X, (0, 1))
+    res = array_derive(cg, X)
+    assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
+
+    cg = ArrayDiagonal(X, (0, 1))
+    res = array_derive(cg, X)
+    assert res == ArrayDiagonal(ArrayTensorProduct(I, I), (1, 3))
+
+    cg = ElementwiseApplyFunction(sin, X)
+    res = array_derive(cg, X)
+    assert res.dummy_eq(ArrayDiagonal(
+        ArrayTensorProduct(
+            ElementwiseApplyFunction(cos, X),
+            I,
+            I
+        ), (0, 3), (1, 5)))
+
+    cg = ArrayElementwiseApplyFunc(sin, X)
+    res = array_derive(cg, X)
+    assert res.dummy_eq(ArrayDiagonal(
+        ArrayTensorProduct(
+            I,
+            I,
+            ArrayElementwiseApplyFunc(cos, X)
+        ), (1, 4), (3, 5)))
+
+    res = array_derive(A1, A1)
+    assert res == PermuteDims(
+        ArrayTensorProduct(Identity(3), Identity(2), Identity(k)),
+        [0, 2, 4, 1, 3, 5]
+    )
+
+    cg = ArrayElementwiseApplyFunc(sin, A1)
+    res = array_derive(cg, A1)
+    assert res.dummy_eq(ArrayDiagonal(
+        ArrayTensorProduct(
+            Identity(3), Identity(2), Identity(k),
+            ArrayElementwiseApplyFunc(cos, A1)
+        ), (1, 6), (3, 7), (5, 8)
+    ))
+
+    cg = Reshape(A, (k**2,))
+    res = array_derive(cg, A)
+    assert res == Reshape(PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3]), (k, k, k**2))

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_as_explicit.py

@@ -0,0 +1,63 @@
+from sympy.core.symbol import Symbol
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensordiagonal, tensorproduct)
+from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
+from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, ArraySymbol, \
+    ArrayTensorProduct, PermuteDims, ArrayDiagonal, ArrayContraction, ArrayAdd
+from sympy.testing.pytest import raises
+
+
+def test_array_as_explicit_call():
+
+    assert ZeroArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray.zeros(3, 2, 4)
+    assert OneArray(3, 2, 4).as_explicit() == ImmutableDenseNDimArray([1 for i in range(3*2*4)]).reshape(3, 2, 4)
+
+    k = Symbol("k")
+    X = ArraySymbol("X", (k, 3, 2))
+    raises(ValueError, lambda: X.as_explicit())
+    raises(ValueError, lambda: ZeroArray(k, 2, 3).as_explicit())
+    raises(ValueError, lambda: OneArray(2, k, 2).as_explicit())
+
+    A = ArraySymbol("A", (3, 3))
+    B = ArraySymbol("B", (3, 3))
+
+    texpr = tensorproduct(A, B)
+    assert isinstance(texpr, ArrayTensorProduct)
+    assert texpr.as_explicit() == tensorproduct(A.as_explicit(), B.as_explicit())
+
+    texpr = tensorcontraction(A, (0, 1))
+    assert isinstance(texpr, ArrayContraction)
+    assert texpr.as_explicit() == A[0, 0] + A[1, 1] + A[2, 2]
+
+    texpr = tensordiagonal(A, (0, 1))
+    assert isinstance(texpr, ArrayDiagonal)
+    assert texpr.as_explicit() == ImmutableDenseNDimArray([A[0, 0], A[1, 1], A[2, 2]])
+
+    texpr = permutedims(A, [1, 0])
+    assert isinstance(texpr, PermuteDims)
+    assert texpr.as_explicit() == permutedims(A.as_explicit(), [1, 0])
+
+
+def test_array_as_explicit_matrix_symbol():
+
+    A = MatrixSymbol("A", 3, 3)
+    B = MatrixSymbol("B", 3, 3)
+
+    texpr = tensorproduct(A, B)
+    assert isinstance(texpr, ArrayTensorProduct)
+    assert texpr.as_explicit() == tensorproduct(A.as_explicit(), B.as_explicit())
+
+    texpr = tensorcontraction(A, (0, 1))
+    assert isinstance(texpr, ArrayContraction)
+    assert texpr.as_explicit() == A[0, 0] + A[1, 1] + A[2, 2]
+
+    texpr = tensordiagonal(A, (0, 1))
+    assert isinstance(texpr, ArrayDiagonal)
+    assert texpr.as_explicit() == ImmutableDenseNDimArray([A[0, 0], A[1, 1], A[2, 2]])
+
+    texpr = permutedims(A, [1, 0])
+    assert isinstance(texpr, PermuteDims)
+    assert texpr.as_explicit() == permutedims(A.as_explicit(), [1, 0])
+
+    expr = ArrayAdd(ArrayTensorProduct(A, B), ArrayTensorProduct(B, A))
+    assert expr.as_explicit() == expr.args[0].as_explicit() + expr.args[1].as_explicit()

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_indexed.py

@@ -0,0 +1,61 @@
+from sympy import Sum, Dummy, sin
+from sympy.tensor.array.expressions import ArraySymbol, ArrayTensorProduct, ArrayContraction, PermuteDims, \
+    ArrayDiagonal, ArrayAdd, OneArray, ZeroArray, convert_indexed_to_array, ArrayElementwiseApplyFunc, Reshape
+from sympy.tensor.array.expressions.from_array_to_indexed import convert_array_to_indexed
+
+from sympy.abc import i, j, k, l, m, n, o
+
+
+def test_convert_array_to_indexed_main():
+    A = ArraySymbol("A", (3, 3, 3))
+    B = ArraySymbol("B", (3, 3))
+    C = ArraySymbol("C", (3, 3))
+
+    d_ = Dummy("d_")
+
+    assert convert_array_to_indexed(A, [i, j, k]) == A[i, j, k]
+
+    expr = ArrayTensorProduct(A, B, C)
+    conv = convert_array_to_indexed(expr, [i,j,k,l,m,n,o])
+    assert conv == A[i,j,k]*B[l,m]*C[n,o]
+    assert convert_indexed_to_array(conv, [i,j,k,l,m,n,o]) == expr
+
+    expr = ArrayContraction(A, (0, 2))
+    assert convert_array_to_indexed(expr, [i]).dummy_eq(Sum(A[d_, i, d_], (d_, 0, 2)))
+
+    expr = ArrayDiagonal(A, (0, 2))
+    assert convert_array_to_indexed(expr, [i, j]) == A[j, i, j]
+
+    expr = PermuteDims(A, [1, 2, 0])
+    conv = convert_array_to_indexed(expr, [i, j, k])
+    assert conv == A[k, i, j]
+    assert convert_indexed_to_array(conv, [i, j, k]) == expr
+
+    expr = ArrayAdd(B, C, PermuteDims(C, [1, 0]))
+    conv = convert_array_to_indexed(expr, [i, j])
+    assert conv == B[i, j] + C[i, j] + C[j, i]
+    assert convert_indexed_to_array(conv, [i, j]) == expr
+
+    expr = ArrayElementwiseApplyFunc(sin, A)
+    conv = convert_array_to_indexed(expr, [i, j, k])
+    assert conv == sin(A[i, j, k])
+    assert convert_indexed_to_array(conv, [i, j, k]).dummy_eq(expr)
+
+    assert convert_array_to_indexed(OneArray(3, 3), [i, j]) == 1
+    assert convert_array_to_indexed(ZeroArray(3, 3), [i, j]) == 0
+
+    expr = Reshape(A, (27,))
+    assert convert_array_to_indexed(expr, [i]) == A[i // 9, i // 3 % 3, i % 3]
+
+    X = ArraySymbol("X", (2, 3, 4, 5, 6))
+    expr = Reshape(X, (2*3*4*5*6,))
+    assert convert_array_to_indexed(expr, [i]) == X[i // 360, i // 120 % 3, i // 30 % 4, i // 6 % 5, i % 6]
+
+    expr = Reshape(X, (4, 9, 2, 2, 5))
+    one_index = 180*i + 20*j + 10*k + 5*l + m
+    expected = X[one_index // (3*4*5*6), one_index // (4*5*6) % 3, one_index // (5*6) % 4, one_index // 6 % 5, one_index % 6]
+    assert convert_array_to_indexed(expr, [i, j, k, l, m]) == expected
+
+    X = ArraySymbol("X", (2*3*5,))
+    expr = Reshape(X, (2, 3, 5))
+    assert convert_array_to_indexed(expr, [i, j, k]) == X[15*i + 5*j + k]

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_matrix.py

@@ -0,0 +1,689 @@
+from sympy import Lambda, S, Dummy, KroneckerProduct
+from sympy.core.symbol import symbols
+from sympy.functions.elementary.miscellaneous import sqrt
+from sympy.functions.elementary.trigonometric import cos, sin
+from sympy.matrices.expressions.hadamard import HadamardProduct, HadamardPower
+from sympy.matrices.expressions.special import (Identity, OneMatrix, ZeroMatrix)
+from sympy.matrices.expressions.matexpr import MatrixElement
+from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+from sympy.tensor.array.expressions.from_array_to_matrix import _support_function_tp1_recognize, \
+    _array_diag2contr_diagmatrix, convert_array_to_matrix, _remove_trivial_dims, _array2matrix, \
+    _combine_removed, identify_removable_identity_matrices, _array_contraction_to_diagonal_multiple_identity
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.combinatorics import Permutation
+from sympy.matrices.expressions.diagonal import DiagMatrix, DiagonalMatrix
+from sympy.matrices import Trace, MatMul, Transpose
+from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, \
+    ArrayElement, ArraySymbol, ArrayElementwiseApplyFunc, _array_tensor_product, _array_contraction, \
+    _array_diagonal, _permute_dims, PermuteDims, ArrayAdd, ArrayDiagonal, ArrayContraction, ArrayTensorProduct
+from sympy.testing.pytest import raises
+
+
+i, j, k, l, m, n = symbols("i j k l m n")
+
+I = Identity(k)
+I1 = Identity(1)
+
+M = MatrixSymbol("M", k, k)
+N = MatrixSymbol("N", k, k)
+P = MatrixSymbol("P", k, k)
+Q = MatrixSymbol("Q", k, k)
+
+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)
+
+a = MatrixSymbol("a", k, 1)
+b = MatrixSymbol("b", k, 1)
+c = MatrixSymbol("c", k, 1)
+d = MatrixSymbol("d", k, 1)
+
+x = MatrixSymbol("x", k, 1)
+y = MatrixSymbol("y", k, 1)
+
+
+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_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_diagonalized_vector():
+
+    # Check matrix recognition over trivial dimensions:
+
+    cg = _array_tensor_product(a, b)
+    assert convert_array_to_matrix(cg) == a * b.T
+
+    cg = _array_tensor_product(I1, a, b)
+    assert convert_array_to_matrix(cg) == a * b.T
+
+    # Recognize trace inside a tensor product:
+
+    cg = _array_contraction(_array_tensor_product(A, B, C), (0, 3), (1, 2))
+    assert convert_array_to_matrix(cg) == Trace(A * B) * C
+
+    # Transform diagonal operator to contraction:
+
+    cg = _array_diagonal(_array_tensor_product(A, a), (1, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (1, 3))
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(a, b), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _permute_dims(
+        _array_contraction(_array_tensor_product(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
+    )
+    assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(A, a), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (0, 3))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(I, x, I1), (0, 2), (3, 5))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
+    assert convert_array_to_matrix(cg) == DiagMatrix(x)
+
+    cg = _array_diagonal(_array_tensor_product(I, x, A, B), (1, 2), (5, 6))
+    assert _array_diag2contr_diagmatrix(cg) == _array_diagonal(_array_contraction(_array_tensor_product(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
+    # TODO: this is returning a wrong result:
+    # convert_array_to_matrix(cg)
+
+    cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3, 5))
+    assert convert_array_to_matrix(cg) == a*b.T
+
+    cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), a, b, I1), (2, 6))
+    assert convert_array_to_matrix(cg) == a*b.T
+
+    cg = _array_diagonal(_array_tensor_product(x, I1), (1, 2))
+    assert isinstance(cg, ArrayDiagonal)
+    assert cg.diagonal_indices == ((1, 2),)
+    assert convert_array_to_matrix(cg) == x
+
+    cg = _array_diagonal(_array_tensor_product(x, I), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), I, DiagMatrix(x)), (1, 3))
+    assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
+
+    raises(ValueError, lambda: _array_diagonal(x, (1,)))
+
+    # Ignore identity matrices with contractions:
+
+    cg = _array_contraction(_array_tensor_product(I, A, I, I), (0, 2), (1, 3), (5, 7))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg) == Trace(A) * I
+
+    cg = _array_contraction(_array_tensor_product(Trace(A) * I, I, I), (1, 5), (3, 4))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg).doit() == Trace(A) * I
+
+    # Add DiagMatrix when required:
+
+    cg = _array_contraction(_array_tensor_product(A, a), (1, 2))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg) == A * a
+
+    cg = _array_contraction(_array_tensor_product(A, a, B), (1, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(A, a, B), (0, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(A, a, b, a.T, B), (0, 2, 4, 7, 9))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1),
+                                                DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
+                                               (0, 2), (3, 5), (6, 9), (8, 12))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
+
+    cg = _array_contraction(_array_tensor_product(I1, I1, I1), (1, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == 1
+
+    cg = _array_contraction(_array_tensor_product(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
+    assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
+
+    cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
+    expected = _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
+    assert cg.split_multiple_contractions() == expected
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
+    expected = _array_contraction(_array_tensor_product(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
+                                (1, 3), (2, 10), (6, 8), (7, 11))
+    assert cg.split_multiple_contractions().dummy_eq(expected)
+    assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
+
+
+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 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_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_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_convert_array_to_matrix_support_function():
+
+    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]) == _array_tensor_product(A * B, C * D)
+    assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
+        _array_tensor_product(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]) == _array_tensor_product(X * Y * A * B, C * D)
+
+    assert _support_function_tp1_recognize([(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims(
+        _array_tensor_product(X * B.T, Y * C, A.T * D.T), [0, 2, 4, 1, 3, 5]
+    )
+
+    assert _support_function_tp1_recognize([(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims(
+        _array_tensor_product(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]) == PermuteDims(
+        _array_tensor_product(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]) == PermuteDims(
+        _array_tensor_product(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_convert_array_to_hadamard_products():
+
+    expr = HadamardProduct(M, N)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = HadamardProduct(M, N)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = Q*HadamardProduct(M, N)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = Q*HadamardProduct(M, N.T)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = HadamardProduct(M, N)*HadamardProduct(Q, P)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert expr == ret
+
+    expr = P.T*HadamardProduct(M, N)*HadamardProduct(Q, P)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert expr == ret
+
+    # ArrayDiagonal should be converted
+    cg = _array_diagonal(_array_tensor_product(M, N, Q), (1, 3), (0, 2, 4))
+    ret = convert_array_to_matrix(cg)
+    expected = PermuteDims(_array_diagonal(_array_tensor_product(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2])
+    assert expected == ret
+
+    # Special case that should return the same expression:
+    cg = _array_diagonal(_array_tensor_product(HadamardProduct(M, N), Q), (0, 2))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    # Hadamard products with traces:
+
+    expr = Trace(HadamardProduct(M, N))
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M.T, N.T))
+
+    expr = Trace(A*HadamardProduct(M, N))
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M, N)*A)
+
+    expr = Trace(HadamardProduct(A, M)*N)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M.T, N)*A)
+
+    # These should not be converted into Hadamard products:
+
+    cg = _array_diagonal(_array_tensor_product(M, N), (0, 1, 2, 3))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    cg = _array_diagonal(_array_tensor_product(A), (0, 1))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 2, 4), (1, 3, 5))
+    assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
+
+    cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 3, 4), (1, 2, 5))
+    assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
+
+    cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
+    assert convert_array_to_matrix(cg) == DiagMatrix(x)
+
+
+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_combine_removed():
+
+    assert _combine_removed(6, [0, 1, 2], [0, 1, 2]) == [0, 1, 2, 3, 4, 5]
+    assert _combine_removed(8, [2, 5], [1, 3, 4]) == [1, 2, 4, 5, 6]
+    assert _combine_removed(8, [7], []) == [7]
+
+
+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_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 test_convert_array_elementwise_function_to_matrix():
+
+    d = Dummy("d")
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x.T*y)
+    assert convert_array_to_matrix(expr) == sin(x.T*y)
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, d**2), x.T*y)
+    assert convert_array_to_matrix(expr) == (x.T*y)**2
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x)
+    assert convert_array_to_matrix(expr).dummy_eq(x.applyfunc(sin))
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, 1 / (2 * sqrt(d))), x)
+    assert convert_array_to_matrix(expr) == S.Half * HadamardPower(x, -S.Half)
+
+
+def test_array2matrix():
+    # See issue https://github.com/sympy/sympy/pull/22877
+    expr = PermuteDims(ArrayContraction(ArrayTensorProduct(x, I, I1, x), (0, 3), (1, 7)), Permutation(2, 3))
+    expected = PermuteDims(ArrayTensorProduct(x*x.T, I1), Permutation(3)(1, 2))
+    assert _array2matrix(expr) == expected
+
+
+def test_recognize_broadcasting():
+    expr = ArrayTensorProduct(x.T*x, A)
+    assert _remove_trivial_dims(expr) == (KroneckerProduct(x.T*x, A), [0, 1])
+
+    expr = ArrayTensorProduct(A, x.T*x)
+    assert _remove_trivial_dims(expr) == (KroneckerProduct(A, x.T*x), [2, 3])
+
+    expr = ArrayTensorProduct(A, B, x.T*x, C)
+    assert _remove_trivial_dims(expr) == (ArrayTensorProduct(A, KroneckerProduct(B, x.T*x), C), [4, 5])
+
+    # Always prefer matrix multiplication to Kronecker product, if possible:
+    expr = ArrayTensorProduct(a, b, x.T*x)
+    assert _remove_trivial_dims(expr) == (a*x.T*x*b.T, [1, 3, 4, 5])

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_convert_indexed_to_array.py

@@ -0,0 +1,205 @@
+from sympy import tanh
+from sympy.concrete.summations import Sum
+from sympy.core.symbol import symbols
+from sympy.functions.special.tensor_functions import KroneckerDelta
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.matrices.expressions.special import Identity
+from sympy.tensor.array.expressions import ArrayElementwiseApplyFunc
+from sympy.tensor.indexed import IndexedBase
+from sympy.combinatorics import Permutation
+from sympy.tensor.array.expressions.array_expressions import ArrayContraction, ArrayTensorProduct, \
+    ArrayDiagonal, ArrayAdd, PermuteDims, ArrayElement, _array_tensor_product, _array_contraction, _array_diagonal, \
+    _array_add, _permute_dims, ArraySymbol, OneArray
+from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array, _convert_indexed_to_array
+from sympy.testing.pytest import raises
+
+
+A, B = symbols("A B", cls=IndexedBase)
+i, j, k, l, m, n = symbols("i j k l m n")
+d0, d1, d2, d3 = symbols("d0:4")
+
+I = Identity(k)
+
+M = MatrixSymbol("M", k, k)
+N = MatrixSymbol("N", k, k)
+P = MatrixSymbol("P", k, k)
+Q = MatrixSymbol("Q", k, k)
+
+a = MatrixSymbol("a", k, 1)
+b = MatrixSymbol("b", k, 1)
+c = MatrixSymbol("c", k, 1)
+d = MatrixSymbol("d", k, 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_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_convert_array_element_to_array_expression():
+    A = ArraySymbol("A", (k,))
+    B = ArraySymbol("B", (k,))
+
+    s = Sum(A[i]*B[i], (i, 0, k-1))
+    cg = convert_indexed_to_array(s)
+    assert cg == ArrayContraction(ArrayTensorProduct(A, B), (0, 1))
+
+    s = A[i]*B[i]
+    cg = convert_indexed_to_array(s)
+    assert cg == ArrayDiagonal(ArrayTensorProduct(A, B), (0, 1))
+
+    s = A[i]*B[j]
+    cg = convert_indexed_to_array(s, [i, j])
+    assert cg == ArrayTensorProduct(A, B)
+    cg = convert_indexed_to_array(s, [j, i])
+    assert cg == ArrayTensorProduct(B, A)
+
+    s = tanh(A[i]*B[j])
+    cg = convert_indexed_to_array(s, [i, j])
+    assert cg.dummy_eq(ArrayElementwiseApplyFunc(tanh, ArrayTensorProduct(A, B)))
+
+
+def test_arrayexpr_convert_indexed_to_array_and_back_to_matrix():
+
+    expr = a.T*b
+    elem = expr[0, 0]
+    cg = convert_indexed_to_array(elem)
+    assert cg == ArrayElement(ArrayContraction(ArrayTensorProduct(a, b), (0, 2)), [0, 0])
+
+    expr = M[i,j] + N[i,j]
+    p1, p2 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(p1) == M + N
+
+    expr = M[i,j] + N[j,i]
+    p1, p2 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(p1) == M + N.T
+
+    expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
+    p1, p2 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(p1) == ArrayAdd(
+        ArrayTensorProduct(M, N),
+        ArrayTensorProduct(N, M))
+
+    expr = (M*N*P)[i, j]
+    p1, p2 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(p1) == M * N * P
+
+    expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
+    p1, p2 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(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 = _convert_indexed_to_array(expr)
+    assert convert_array_to_matrix(p1) == M * P * N + M * P.T * N + N * P * M + N * P.T * M
+
+
+def test_arrayexpr_convert_indexed_to_array_out_of_bounds():
+
+    expr = Sum(M[i, i], (i, 0, 4))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+    expr = Sum(M[i, i], (i, 0, k))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+    expr = Sum(M[i, i], (i, 1, k-1))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+
+    expr = Sum(M[i, j]*N[j,m], (j, 0, 4))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+    expr = Sum(M[i, j]*N[j,m], (j, 0, k))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+    expr = Sum(M[i, j]*N[j,m], (j, 1, k-1))
+    raises(ValueError, lambda: convert_indexed_to_array(expr))
+
+
+def test_arrayexpr_convert_indexed_to_array_broadcast():
+    A = ArraySymbol("A", (3, 3))
+    B = ArraySymbol("B", (3, 3))
+
+    expr = A[i, j] + B[k, l]
+    O2 = OneArray(3, 3)
+    expected = ArrayAdd(ArrayTensorProduct(A, O2), ArrayTensorProduct(O2, B))
+    assert convert_indexed_to_array(expr) == expected
+    assert convert_indexed_to_array(expr, [i, j, k, l]) == expected
+    assert convert_indexed_to_array(expr, [l, k, i, j]) == ArrayAdd(PermuteDims(ArrayTensorProduct(O2, A), [1, 0, 2, 3]), PermuteDims(ArrayTensorProduct(B, O2), [1, 0, 2, 3]))
+
+    expr = A[i, j] + B[j, k]
+    O1 = OneArray(3)
+    assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(ArrayTensorProduct(A, O1), ArrayTensorProduct(O1, B))
+
+    C = ArraySymbol("C", (d0, d1))
+    D = ArraySymbol("D", (d3, d1))
+
+    expr = C[i, j] + D[k, j]
+    assert convert_indexed_to_array(expr, [i, j, k]) == ArrayAdd(ArrayTensorProduct(C, OneArray(d3)), PermuteDims(ArrayTensorProduct(OneArray(d0), D), [0, 2, 1]))
+
+    X = ArraySymbol("X", (5, 3))
+
+    expr = X[i, n] - X[j, n]
+    assert convert_indexed_to_array(expr, [i, j, n]) == ArrayAdd(ArrayTensorProduct(-1, OneArray(5), X), PermuteDims(ArrayTensorProduct(X, OneArray(5)), [0, 2, 1]))
+
+    raises(ValueError, lambda: convert_indexed_to_array(C[i, j] + D[i, j]))

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_convert_matrix_to_array.py

@@ -0,0 +1,128 @@
+from sympy import Lambda, KroneckerProduct
+from sympy.core.symbol import symbols, Dummy
+from sympy.matrices.expressions.hadamard import (HadamardPower, HadamardProduct)
+from sympy.matrices.expressions.inverse import Inverse
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.matrices.expressions.matpow import MatPow
+from sympy.matrices.expressions.special import Identity
+from sympy.matrices.expressions.trace import Trace
+from sympy.matrices.expressions.transpose import Transpose
+from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayContraction, \
+    PermuteDims, ArrayDiagonal, ArrayElementwiseApplyFunc, _array_contraction, _array_tensor_product, Reshape
+from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+
+i, j, k, l, m, n = symbols("i j k l m n")
+
+I = Identity(k)
+
+M = MatrixSymbol("M", k, k)
+N = MatrixSymbol("N", k, k)
+P = MatrixSymbol("P", k, k)
+Q = MatrixSymbol("Q", k, k)
+
+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)
+
+a = MatrixSymbol("a", k, 1)
+b = MatrixSymbol("b", k, 1)
+c = MatrixSymbol("c", k, 1)
+d = MatrixSymbol("d", k, 1)
+
+
+def test_arrayexpr_convert_matrix_to_array():
+
+    expr = M*N
+    result = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = M*N*M
+    result = _array_contraction(ArrayTensorProduct(M, N, M), (1, 2), (3, 4))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = Transpose(M)
+    assert convert_matrix_to_array(expr) == PermuteDims(M, [1, 0])
+
+    expr = M*Transpose(N)
+    assert convert_matrix_to_array(expr) == _array_contraction(_array_tensor_product(M, PermuteDims(N, [1, 0])), (1, 2))
+
+    expr = 3*M*N
+    res = convert_matrix_to_array(expr)
+    rexpr = convert_array_to_matrix(res)
+    assert expr == rexpr
+
+    expr = 3*M + N*M.T*M + 4*k*N
+    res = convert_matrix_to_array(expr)
+    rexpr = convert_array_to_matrix(res)
+    assert expr == rexpr
+
+    expr = Inverse(M)*N
+    rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
+    assert expr == rexpr
+
+    expr = M**2
+    rexpr = convert_array_to_matrix(convert_matrix_to_array(expr))
+    assert expr == rexpr
+
+    expr = M*(2*N + 3*M)
+    res = convert_matrix_to_array(expr)
+    rexpr = convert_array_to_matrix(res)
+    assert expr == rexpr
+
+    expr = Trace(M)
+    result = ArrayContraction(M, (0, 1))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = 3*Trace(M)
+    result = ArrayContraction(ArrayTensorProduct(3, M), (0, 1))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = 3*Trace(Trace(M) * M)
+    result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = 3*Trace(M)**2
+    result = ArrayContraction(ArrayTensorProduct(3, M, M), (0, 1), (2, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = HadamardProduct(M, N)
+    result = ArrayDiagonal(ArrayTensorProduct(M, N), (0, 2), (1, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = HadamardProduct(M*N, N*M)
+    result = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, N, M), (1, 2), (5, 6)), (0, 2), (1, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = HadamardPower(M, 2)
+    result = ArrayDiagonal(ArrayTensorProduct(M, M), (0, 2), (1, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = HadamardPower(M*N, 2)
+    result = ArrayDiagonal(ArrayContraction(ArrayTensorProduct(M, N, M, N), (1, 2), (5, 6)), (0, 2), (1, 3))
+    assert convert_matrix_to_array(expr) == result
+
+    expr = HadamardPower(M, n)
+    d0 = Dummy("d0")
+    result = ArrayElementwiseApplyFunc(Lambda(d0, d0**n), M)
+    assert convert_matrix_to_array(expr).dummy_eq(result)
+
+    expr = M**2
+    assert isinstance(expr, MatPow)
+    assert convert_matrix_to_array(expr) == ArrayContraction(ArrayTensorProduct(M, M), (1, 2))
+
+    expr = a.T*b
+    cg = convert_matrix_to_array(expr)
+    assert cg == ArrayContraction(ArrayTensorProduct(a, b), (0, 2))
+
+    expr = KroneckerProduct(A, B)
+    cg = convert_matrix_to_array(expr)
+    assert cg == Reshape(PermuteDims(ArrayTensorProduct(A, B), [0, 2, 1, 3]), (k**2, k**2))
+
+    expr = KroneckerProduct(A, B, C, D)
+    cg = convert_matrix_to_array(expr)
+    assert cg == Reshape(PermuteDims(ArrayTensorProduct(A, B, C, D), [0, 2, 4, 6, 1, 3, 5, 7]), (k**4, k**4))

rl/Lib/site-packages/sympy/tensor/array/expressions/tests/test_deprecated_conv_modules.py

@@ -0,0 +1,22 @@
+from sympy import MatrixSymbol, symbols, Sum
+from sympy.tensor.array.expressions import conv_array_to_indexed, from_array_to_indexed, ArrayTensorProduct, \
+    ArrayContraction, conv_array_to_matrix, from_array_to_matrix, conv_matrix_to_array, from_matrix_to_array, \
+    conv_indexed_to_array, from_indexed_to_array
+from sympy.testing.pytest import warns
+from sympy.utilities.exceptions import SymPyDeprecationWarning
+
+
+def test_deprecated_conv_module_results():
+
+    M = MatrixSymbol("M", 3, 3)
+    N = MatrixSymbol("N", 3, 3)
+    i, j, d = symbols("i j d")
+
+    x = ArrayContraction(ArrayTensorProduct(M, N), (1, 2))
+    y = Sum(M[i, d]*N[d, j], (d, 0, 2))
+
+    with warns(SymPyDeprecationWarning, test_stacklevel=False):
+        assert conv_array_to_indexed.convert_array_to_indexed(x, [i, j]).dummy_eq(from_array_to_indexed.convert_array_to_indexed(x, [i, j]))
+        assert conv_array_to_matrix.convert_array_to_matrix(x) == from_array_to_matrix.convert_array_to_matrix(x)
+        assert conv_matrix_to_array.convert_matrix_to_array(M*N) == from_matrix_to_array.convert_matrix_to_array(M*N)
+        assert conv_indexed_to_array.convert_indexed_to_array(y) == from_indexed_to_array.convert_indexed_to_array(y)

rl/Lib/site-packages/sympy/tensor/array/expressions/utils.py

@@ -0,0 +1,123 @@
+import bisect
+from collections import defaultdict
+
+from sympy.combinatorics import Permutation
+from sympy.core.containers import Tuple
+from sympy.core.numbers import Integer
+
+
+def _get_mapping_from_subranks(subranks):
+    mapping = {}
+    counter = 0
+    for i, rank in enumerate(subranks):
+        for j in range(rank):
+            mapping[counter] = (i, j)
+            counter += 1
+    return mapping
+
+
+def _get_contraction_links(args, subranks, *contraction_indices):
+    mapping = _get_mapping_from_subranks(subranks)
+    contraction_tuples = [[mapping[j] for j in i] for i in contraction_indices]
+    dlinks = defaultdict(dict)
+    for links in contraction_tuples:
+        if len(links) == 2:
+            (arg1, pos1), (arg2, pos2) = links
+            dlinks[arg1][pos1] = (arg2, pos2)
+            dlinks[arg2][pos2] = (arg1, pos1)
+            continue
+
+    return args, dict(dlinks)
+
+
+def _sort_contraction_indices(pairing_indices):
+    pairing_indices = [Tuple(*sorted(i)) for i in pairing_indices]
+    pairing_indices.sort(key=lambda x: min(x))
+    return pairing_indices
+
+
+def _get_diagonal_indices(flattened_indices):
+    axes_contraction = defaultdict(list)
+    for i, ind in enumerate(flattened_indices):
+        if isinstance(ind, (int, Integer)):
+            # If the indices is a number, there can be no diagonal operation:
+            continue
+        axes_contraction[ind].append(i)
+    axes_contraction = {k: v for k, v in axes_contraction.items() if len(v) > 1}
+    # Put the diagonalized indices at the end:
+    ret_indices = [i for i in flattened_indices if i not in axes_contraction]
+    diag_indices = list(axes_contraction)
+    diag_indices.sort(key=lambda x: flattened_indices.index(x))
+    diagonal_indices = [tuple(axes_contraction[i]) for i in diag_indices]
+    ret_indices += diag_indices
+    ret_indices = tuple(ret_indices)
+    return diagonal_indices, ret_indices
+
+
+def _get_argindex(subindices, ind):
+    for i, sind in enumerate(subindices):
+        if ind == sind:
+            return i
+        if isinstance(sind, (set, frozenset)) and ind in sind:
+            return i
+    raise IndexError("%s not found in %s" % (ind, subindices))
+
+
+def _apply_recursively_over_nested_lists(func, arr):
+    if isinstance(arr, (tuple, list, Tuple)):
+        return tuple(_apply_recursively_over_nested_lists(func, i) for i in arr)
+    elif isinstance(arr, Tuple):
+        return Tuple.fromiter(_apply_recursively_over_nested_lists(func, i) for i in arr)
+    else:
+        return func(arr)
+
+
+def _build_push_indices_up_func_transformation(flattened_contraction_indices):
+    shifts = {0: 0}
+    i = 0
+    cumulative = 0
+    while i < len(flattened_contraction_indices):
+        j = 1
+        while i+j < len(flattened_contraction_indices):
+            if flattened_contraction_indices[i] + j != flattened_contraction_indices[i+j]:
+                break
+            j += 1
+        cumulative += j
+        shifts[flattened_contraction_indices[i]] = cumulative
+        i += j
+    shift_keys = sorted(shifts.keys())
+
+    def func(idx):
+        return shifts[shift_keys[bisect.bisect_right(shift_keys, idx)-1]]
+
+    def transform(j):
+        if j in flattened_contraction_indices:
+            return None
+        else:
+            return j - func(j)
+
+    return transform
+
+
+def _build_push_indices_down_func_transformation(flattened_contraction_indices):
+    N = flattened_contraction_indices[-1]+2
+
+    shifts = [i for i in range(N) if i not in flattened_contraction_indices]
+
+    def transform(j):
+        if j < len(shifts):
+            return shifts[j]
+        else:
+            return j + shifts[-1] - len(shifts) + 1
+
+    return transform
+
+
+def _apply_permutation_to_list(perm: Permutation, target_list: list):
+    """
+    Permute a list according to the given permutation.
+    """
+    new_list = [None for i in range(perm.size)]
+    for i, e in enumerate(target_list):
+        new_list[perm(i)] = e
+    return new_list