Add files using upload-large-folder tool

.gitattributes

@@ -437,3 +437,5 @@ tuning-competition-baseline/.venv/lib/python3.11/site-packages/nvidia/cudnn/lib/
 .venv/lib/python3.11/site-packages/transformers/models/wav2vec2/__pycache__/modeling_tf_wav2vec2.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
 .venv/lib/python3.11/site-packages/transformers/models/wav2vec2_conformer/__pycache__/modeling_wav2vec2_conformer.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
 .venv/lib/python3.11/site-packages/transformers/models/whisper/__pycache__/modeling_whisper.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_tf_objects.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+.venv/lib/python3.11/site-packages/transformers/utils/__pycache__/dummy_pt_objects.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text

.venv/lib/python3.11/site-packages/sympy/tensor/__init__.py

@@ -0,0 +1,23 @@
+"""A module to manipulate symbolic objects with indices including tensors
+
+"""
+from .indexed import IndexedBase, Idx, Indexed
+from .index_methods import get_contraction_structure, get_indices
+from .functions import shape
+from .array import (MutableDenseNDimArray, ImmutableDenseNDimArray,
+    MutableSparseNDimArray, ImmutableSparseNDimArray, NDimArray, tensorproduct,
+    tensorcontraction, tensordiagonal, derive_by_array, permutedims, Array,
+    DenseNDimArray, SparseNDimArray,)
+
+__all__ = [
+    'IndexedBase', 'Idx', 'Indexed',
+
+    'get_contraction_structure', 'get_indices',
+
+    'shape',
+
+    'MutableDenseNDimArray', 'ImmutableDenseNDimArray',
+    'MutableSparseNDimArray', 'ImmutableSparseNDimArray', 'NDimArray',
+    'tensorproduct', 'tensorcontraction', 'tensordiagonal', 'derive_by_array', 'permutedims',
+    'Array', 'DenseNDimArray', 'SparseNDimArray',
+]

.venv/lib/python3.11/site-packages/sympy/tensor/array/__init__.py

@@ -0,0 +1,271 @@
+r"""
+N-dim array module for SymPy.
+
+Four classes are provided to handle N-dim arrays, given by the combinations
+dense/sparse (i.e. whether to store all elements or only the non-zero ones in
+memory) and mutable/immutable (immutable classes are SymPy objects, but cannot
+change after they have been created).
+
+Examples
+========
+
+The following examples show the usage of ``Array``. This is an abbreviation for
+``ImmutableDenseNDimArray``, that is an immutable and dense N-dim array, the
+other classes are analogous. For mutable classes it is also possible to change
+element values after the object has been constructed.
+
+Array construction can detect the shape of nested lists and tuples:
+
+>>> from sympy import Array
+>>> a1 = Array([[1, 2], [3, 4], [5, 6]])
+>>> a1
+[[1, 2], [3, 4], [5, 6]]
+>>> a1.shape
+(3, 2)
+>>> a1.rank()
+2
+>>> from sympy.abc import x, y, z
+>>> a2 = Array([[[x, y], [z, x*z]], [[1, x*y], [1/x, x/y]]])
+>>> a2
+[[[x, y], [z, x*z]], [[1, x*y], [1/x, x/y]]]
+>>> a2.shape
+(2, 2, 2)
+>>> a2.rank()
+3
+
+Otherwise one could pass a 1-dim array followed by a shape tuple:
+
+>>> m1 = Array(range(12), (3, 4))
+>>> m1
+[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
+>>> m2 = Array(range(12), (3, 2, 2))
+>>> m2
+[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8, 9], [10, 11]]]
+>>> m2[1,1,1]
+7
+>>> m2.reshape(4, 3)
+[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
+
+Slice support:
+
+>>> m2[:, 1, 1]
+[3, 7, 11]
+
+Elementwise derivative:
+
+>>> from sympy.abc import x, y, z
+>>> m3 = Array([x**3, x*y, z])
+>>> m3.diff(x)
+[3*x**2, y, 0]
+>>> m3.diff(z)
+[0, 0, 1]
+
+Multiplication with other SymPy expressions is applied elementwisely:
+
+>>> (1+x)*m3
+[x**3*(x + 1), x*y*(x + 1), z*(x + 1)]
+
+To apply a function to each element of the N-dim array, use ``applyfunc``:
+
+>>> m3.applyfunc(lambda x: x/2)
+[x**3/2, x*y/2, z/2]
+
+N-dim arrays can be converted to nested lists by the ``tolist()`` method:
+
+>>> m2.tolist()
+[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8, 9], [10, 11]]]
+>>> isinstance(m2.tolist(), list)
+True
+
+If the rank is 2, it is possible to convert them to matrices with ``tomatrix()``:
+
+>>> m1.tomatrix()
+Matrix([
+[0, 1,  2,  3],
+[4, 5,  6,  7],
+[8, 9, 10, 11]])
+
+Products and contractions
+-------------------------
+
+Tensor product between arrays `A_{i_1,\ldots,i_n}` and `B_{j_1,\ldots,j_m}`
+creates the combined array `P = A \otimes B` defined as
+
+`P_{i_1,\ldots,i_n,j_1,\ldots,j_m} := A_{i_1,\ldots,i_n}\cdot B_{j_1,\ldots,j_m}.`
+
+It is available through ``tensorproduct(...)``:
+
+>>> from sympy import Array, tensorproduct
+>>> from sympy.abc import x,y,z,t
+>>> A = Array([x, y, z, t])
+>>> B = Array([1, 2, 3, 4])
+>>> tensorproduct(A, B)
+[[x, 2*x, 3*x, 4*x], [y, 2*y, 3*y, 4*y], [z, 2*z, 3*z, 4*z], [t, 2*t, 3*t, 4*t]]
+
+In case you don't want to evaluate the tensor product immediately, you can use
+``ArrayTensorProduct``, which creates an unevaluated tensor product expression:
+
+>>> from sympy.tensor.array.expressions import ArrayTensorProduct
+>>> ArrayTensorProduct(A, B)
+ArrayTensorProduct([x, y, z, t], [1, 2, 3, 4])
+
+Calling ``.as_explicit()`` on ``ArrayTensorProduct`` is equivalent to just calling
+``tensorproduct(...)``:
+
+>>> ArrayTensorProduct(A, B).as_explicit()
+[[x, 2*x, 3*x, 4*x], [y, 2*y, 3*y, 4*y], [z, 2*z, 3*z, 4*z], [t, 2*t, 3*t, 4*t]]
+
+Tensor product between a rank-1 array and a matrix creates a rank-3 array:
+
+>>> from sympy import eye
+>>> p1 = tensorproduct(A, eye(4))
+>>> p1
+[[[x, 0, 0, 0], [0, x, 0, 0], [0, 0, x, 0], [0, 0, 0, x]], [[y, 0, 0, 0], [0, y, 0, 0], [0, 0, y, 0], [0, 0, 0, y]], [[z, 0, 0, 0], [0, z, 0, 0], [0, 0, z, 0], [0, 0, 0, z]], [[t, 0, 0, 0], [0, t, 0, 0], [0, 0, t, 0], [0, 0, 0, t]]]
+
+Now, to get back `A_0 \otimes \mathbf{1}` one can access `p_{0,m,n}` by slicing:
+
+>>> p1[0,:,:]
+[[x, 0, 0, 0], [0, x, 0, 0], [0, 0, x, 0], [0, 0, 0, x]]
+
+Tensor contraction sums over the specified axes, for example contracting
+positions `a` and `b` means
+
+`A_{i_1,\ldots,i_a,\ldots,i_b,\ldots,i_n} \implies \sum_k A_{i_1,\ldots,k,\ldots,k,\ldots,i_n}`
+
+Remember that Python indexing is zero starting, to contract the a-th and b-th
+axes it is therefore necessary to specify `a-1` and `b-1`
+
+>>> from sympy import tensorcontraction
+>>> C = Array([[x, y], [z, t]])
+
+The matrix trace is equivalent to the contraction of a rank-2 array:
+
+`A_{m,n} \implies \sum_k A_{k,k}`
+
+>>> tensorcontraction(C, (0, 1))
+t + x
+
+To create an expression representing a tensor contraction that does not get
+evaluated immediately, use ``ArrayContraction``, which is equivalent to
+``tensorcontraction(...)`` if it is followed by ``.as_explicit()``:
+
+>>> from sympy.tensor.array.expressions import ArrayContraction
+>>> ArrayContraction(C, (0, 1))
+ArrayContraction([[x, y], [z, t]], (0, 1))
+>>> ArrayContraction(C, (0, 1)).as_explicit()
+t + x
+
+Matrix product is equivalent to a tensor product of two rank-2 arrays, followed
+by a contraction of the 2nd and 3rd axes (in Python indexing axes number 1, 2).
+
+`A_{m,n}\cdot B_{i,j} \implies \sum_k A_{m, k}\cdot B_{k, j}`
+
+>>> D = Array([[2, 1], [0, -1]])
+>>> tensorcontraction(tensorproduct(C, D), (1, 2))
+[[2*x, x - y], [2*z, -t + z]]
+
+One may verify that the matrix product is equivalent:
+
+>>> from sympy import Matrix
+>>> Matrix([[x, y], [z, t]])*Matrix([[2, 1], [0, -1]])
+Matrix([
+[2*x,  x - y],
+[2*z, -t + z]])
+
+or equivalently
+
+>>> C.tomatrix()*D.tomatrix()
+Matrix([
+[2*x,  x - y],
+[2*z, -t + z]])
+
+Diagonal operator
+-----------------
+
+The ``tensordiagonal`` function acts in a similar manner as ``tensorcontraction``,
+but the joined indices are not summed over, for example diagonalizing
+positions `a` and `b` means
+
+`A_{i_1,\ldots,i_a,\ldots,i_b,\ldots,i_n} \implies A_{i_1,\ldots,k,\ldots,k,\ldots,i_n}
+\implies \tilde{A}_{i_1,\ldots,i_{a-1},i_{a+1},\ldots,i_{b-1},i_{b+1},\ldots,i_n,k}`
+
+where `\tilde{A}` is the array equivalent to the diagonal of `A` at positions
+`a` and `b` moved to the last index slot.
+
+Compare the difference between contraction and diagonal operators:
+
+>>> from sympy import tensordiagonal
+>>> from sympy.abc import a, b, c, d
+>>> m = Matrix([[a, b], [c, d]])
+>>> tensorcontraction(m, [0, 1])
+a + d
+>>> tensordiagonal(m, [0, 1])
+[a, d]
+
+In short, no summation occurs with ``tensordiagonal``.
+
+
+Derivatives by array
+--------------------
+
+The usual derivative operation may be extended to support derivation with
+respect to arrays, provided that all elements in the that array are symbols or
+expressions suitable for derivations.
+
+The definition of a derivative by an array is as follows: given the array
+`A_{i_1, \ldots, i_N}` and the array `X_{j_1, \ldots, j_M}`
+the derivative of arrays will return a new array `B` defined by
+
+`B_{j_1,\ldots,j_M,i_1,\ldots,i_N} := \frac{\partial A_{i_1,\ldots,i_N}}{\partial X_{j_1,\ldots,j_M}}`
+
+The function ``derive_by_array`` performs such an operation:
+
+>>> from sympy import derive_by_array
+>>> from sympy.abc import x, y, z, t
+>>> from sympy import sin, exp
+
+With scalars, it behaves exactly as the ordinary derivative:
+
+>>> derive_by_array(sin(x*y), x)
+y*cos(x*y)
+
+Scalar derived by an array basis:
+
+>>> derive_by_array(sin(x*y), [x, y, z])
+[y*cos(x*y), x*cos(x*y), 0]
+
+Deriving array by an array basis: `B^{nm} := \frac{\partial A^m}{\partial x^n}`
+
+>>> basis = [x, y, z]
+>>> ax = derive_by_array([exp(x), sin(y*z), t], basis)
+>>> ax
+[[exp(x), 0, 0], [0, z*cos(y*z), 0], [0, y*cos(y*z), 0]]
+
+Contraction of the resulting array: `\sum_m \frac{\partial A^m}{\partial x^m}`
+
+>>> tensorcontraction(ax, (0, 1))
+z*cos(y*z) + exp(x)
+
+"""
+
+from .dense_ndim_array import MutableDenseNDimArray, ImmutableDenseNDimArray, DenseNDimArray
+from .sparse_ndim_array import MutableSparseNDimArray, ImmutableSparseNDimArray, SparseNDimArray
+from .ndim_array import NDimArray, ArrayKind
+from .arrayop import tensorproduct, tensorcontraction, tensordiagonal, derive_by_array, permutedims
+from .array_comprehension import ArrayComprehension, ArrayComprehensionMap
+
+Array = ImmutableDenseNDimArray
+
+__all__ = [
+    'MutableDenseNDimArray', 'ImmutableDenseNDimArray', 'DenseNDimArray',
+
+    'MutableSparseNDimArray', 'ImmutableSparseNDimArray', 'SparseNDimArray',
+
+    'NDimArray', 'ArrayKind',
+
+    'tensorproduct', 'tensorcontraction', 'tensordiagonal', 'derive_by_array',
+
+    'permutedims', 'ArrayComprehension', 'ArrayComprehensionMap',
+
+    'Array',
+]

.venv/lib/python3.11/site-packages/sympy/tensor/array/array_comprehension.py

@@ -0,0 +1,399 @@
+import functools, itertools
+from sympy.core.sympify import _sympify, sympify
+from sympy.core.expr import Expr
+from sympy.core import Basic, Tuple
+from sympy.tensor.array import ImmutableDenseNDimArray
+from sympy.core.symbol import Symbol
+from sympy.core.numbers import Integer
+
+
+class ArrayComprehension(Basic):
+    """
+    Generate a list comprehension.
+
+    Explanation
+    ===========
+
+    If there is a symbolic dimension, for example, say [i for i in range(1, N)] where
+    N is a Symbol, then the expression will not be expanded to an array. Otherwise,
+    calling the doit() function will launch the expansion.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array import ArrayComprehension
+    >>> from sympy import symbols
+    >>> i, j, k = symbols('i j k')
+    >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+    >>> a
+    ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+    >>> a.doit()
+    [[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]]
+    >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k))
+    >>> b.doit()
+    ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k))
+    """
+    def __new__(cls, function, *symbols, **assumptions):
+        if any(len(l) != 3 or None for l in symbols):
+            raise ValueError('ArrayComprehension requires values lower and upper bound'
+                              ' for the expression')
+        arglist = [sympify(function)]
+        arglist.extend(cls._check_limits_validity(function, symbols))
+        obj = Basic.__new__(cls, *arglist, **assumptions)
+        obj._limits = obj._args[1:]
+        obj._shape = cls._calculate_shape_from_limits(obj._limits)
+        obj._rank = len(obj._shape)
+        obj._loop_size = cls._calculate_loop_size(obj._shape)
+        return obj
+
+    @property
+    def function(self):
+        """The function applied across limits.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j = symbols('i j')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.function
+        10*i + j
+        """
+        return self._args[0]
+
+    @property
+    def limits(self):
+        """
+        The list of limits that will be applied while expanding the array.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j = symbols('i j')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.limits
+        ((i, 1, 4), (j, 1, 3))
+        """
+        return self._limits
+
+    @property
+    def free_symbols(self):
+        """
+        The set of the free_symbols in the array.
+        Variables appeared in the bounds are supposed to be excluded
+        from the free symbol set.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j, k = symbols('i j k')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.free_symbols
+        set()
+        >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
+        >>> b.free_symbols
+        {k}
+        """
+        expr_free_sym = self.function.free_symbols
+        for var, inf, sup in self._limits:
+            expr_free_sym.discard(var)
+            curr_free_syms = inf.free_symbols.union(sup.free_symbols)
+            expr_free_sym = expr_free_sym.union(curr_free_syms)
+        return expr_free_sym
+
+    @property
+    def variables(self):
+        """The tuples of the variables in the limits.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j, k = symbols('i j k')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.variables
+        [i, j]
+        """
+        return [l[0] for l in self._limits]
+
+    @property
+    def bound_symbols(self):
+        """The list of dummy variables.
+
+        Note
+        ====
+
+        Note that all variables are dummy variables since a limit without
+        lower bound or upper bound is not accepted.
+        """
+        return [l[0] for l in self._limits if len(l) != 1]
+
+    @property
+    def shape(self):
+        """
+        The shape of the expanded array, which may have symbols.
+
+        Note
+        ====
+
+        Both the lower and the upper bounds are included while
+        calculating the shape.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j, k = symbols('i j k')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.shape
+        (4, 3)
+        >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
+        >>> b.shape
+        (4, k + 3)
+        """
+        return self._shape
+
+    @property
+    def is_shape_numeric(self):
+        """
+        Test if the array is shape-numeric which means there is no symbolic
+        dimension.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j, k = symbols('i j k')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.is_shape_numeric
+        True
+        >>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3))
+        >>> b.is_shape_numeric
+        False
+        """
+        for _, inf, sup in self._limits:
+            if Basic(inf, sup).atoms(Symbol):
+                return False
+        return True
+
+    def rank(self):
+        """The rank of the expanded array.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j, k = symbols('i j k')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.rank()
+        2
+        """
+        return self._rank
+
+    def __len__(self):
+        """
+        The length of the expanded array which means the number
+        of elements in the array.
+
+        Raises
+        ======
+
+        ValueError : When the length of the array is symbolic
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j = symbols('i j')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> len(a)
+        12
+        """
+        if self._loop_size.free_symbols:
+            raise ValueError('Symbolic length is not supported')
+        return self._loop_size
+
+    @classmethod
+    def _check_limits_validity(cls, function, limits):
+        #limits = sympify(limits)
+        new_limits = []
+        for var, inf, sup in limits:
+            var = _sympify(var)
+            inf = _sympify(inf)
+            #since this is stored as an argument, it should be
+            #a Tuple
+            if isinstance(sup, list):
+                sup = Tuple(*sup)
+            else:
+                sup = _sympify(sup)
+            new_limits.append(Tuple(var, inf, sup))
+            if any((not isinstance(i, Expr)) or i.atoms(Symbol, Integer) != i.atoms()
+                                                                for i in [inf, sup]):
+                raise TypeError('Bounds should be an Expression(combination of Integer and Symbol)')
+            if (inf > sup) == True:
+                raise ValueError('Lower bound should be inferior to upper bound')
+            if var in inf.free_symbols or var in sup.free_symbols:
+                raise ValueError('Variable should not be part of its bounds')
+        return new_limits
+
+    @classmethod
+    def _calculate_shape_from_limits(cls, limits):
+        return tuple([sup - inf + 1 for _, inf, sup in limits])
+
+    @classmethod
+    def _calculate_loop_size(cls, shape):
+        if not shape:
+            return 0
+        loop_size = 1
+        for l in shape:
+            loop_size = loop_size * l
+
+        return loop_size
+
+    def doit(self, **hints):
+        if not self.is_shape_numeric:
+            return self
+
+        return self._expand_array()
+
+    def _expand_array(self):
+        res = []
+        for values in itertools.product(*[range(inf, sup+1)
+                                        for var, inf, sup
+                                        in self._limits]):
+            res.append(self._get_element(values))
+
+        return ImmutableDenseNDimArray(res, self.shape)
+
+    def _get_element(self, values):
+        temp = self.function
+        for var, val in zip(self.variables, values):
+            temp = temp.subs(var, val)
+        return temp
+
+    def tolist(self):
+        """Transform the expanded array to a list.
+
+        Raises
+        ======
+
+        ValueError : When there is a symbolic dimension
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j = symbols('i j')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.tolist()
+        [[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]]
+        """
+        if self.is_shape_numeric:
+            return self._expand_array().tolist()
+
+        raise ValueError("A symbolic array cannot be expanded to a list")
+
+    def tomatrix(self):
+        """Transform the expanded array to a matrix.
+
+        Raises
+        ======
+
+        ValueError : When there is a symbolic dimension
+        ValueError : When the rank of the expanded array is not equal to 2
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import ArrayComprehension
+        >>> from sympy import symbols
+        >>> i, j = symbols('i j')
+        >>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3))
+        >>> a.tomatrix()
+        Matrix([
+        [11, 12, 13],
+        [21, 22, 23],
+        [31, 32, 33],
+        [41, 42, 43]])
+        """
+        from sympy.matrices import Matrix
+
+        if not self.is_shape_numeric:
+            raise ValueError("A symbolic array cannot be expanded to a matrix")
+        if self._rank != 2:
+            raise ValueError('Dimensions must be of size of 2')
+
+        return Matrix(self._expand_array().tomatrix())
+
+
+def isLambda(v):
+    LAMBDA = lambda: 0
+    return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__
+
+class ArrayComprehensionMap(ArrayComprehension):
+    '''
+    A subclass of ArrayComprehension dedicated to map external function lambda.
+
+    Notes
+    =====
+
+    Only the lambda function is considered.
+    At most one argument in lambda function is accepted in order to avoid ambiguity
+    in value assignment.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array import ArrayComprehensionMap
+    >>> from sympy import symbols
+    >>> i, j, k = symbols('i j k')
+    >>> a = ArrayComprehensionMap(lambda: 1, (i, 1, 4))
+    >>> a.doit()
+    [1, 1, 1, 1]
+    >>> b = ArrayComprehensionMap(lambda a: a+1, (j, 1, 4))
+    >>> b.doit()
+    [2, 3, 4, 5]
+
+    '''
+    def __new__(cls, function, *symbols, **assumptions):
+        if any(len(l) != 3 or None for l in symbols):
+            raise ValueError('ArrayComprehension requires values lower and upper bound'
+                              ' for the expression')
+
+        if not isLambda(function):
+            raise ValueError('Data type not supported')
+
+        arglist = cls._check_limits_validity(function, symbols)
+        obj = Basic.__new__(cls, *arglist, **assumptions)
+        obj._limits = obj._args
+        obj._shape = cls._calculate_shape_from_limits(obj._limits)
+        obj._rank = len(obj._shape)
+        obj._loop_size = cls._calculate_loop_size(obj._shape)
+        obj._lambda = function
+        return obj
+
+    @property
+    def func(self):
+        class _(ArrayComprehensionMap):
+            def __new__(cls, *args, **kwargs):
+                return ArrayComprehensionMap(self._lambda, *args, **kwargs)
+        return _
+
+    def _get_element(self, values):
+        temp = self._lambda
+        if self._lambda.__code__.co_argcount == 0:
+            temp = temp()
+        elif self._lambda.__code__.co_argcount == 1:
+            temp = temp(functools.reduce(lambda a, b: a*b, values))
+        return temp

.venv/lib/python3.11/site-packages/sympy/tensor/array/array_derivatives.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from sympy.core.expr import Expr
+from sympy.core.function import Derivative
+from sympy.core.numbers import Integer
+from sympy.matrices.matrixbase import MatrixBase
+from .ndim_array import NDimArray
+from .arrayop import derive_by_array
+from sympy.matrices.expressions.matexpr import MatrixExpr
+from sympy.matrices.expressions.special import ZeroMatrix
+from sympy.matrices.expressions.matexpr import _matrix_derivative
+
+
+class ArrayDerivative(Derivative):
+
+    is_scalar = False
+
+    def __new__(cls, expr, *variables, **kwargs):
+        obj = super().__new__(cls, expr, *variables, **kwargs)
+        if isinstance(obj, ArrayDerivative):
+            obj._shape = obj._get_shape()
+        return obj
+
+    def _get_shape(self):
+        shape = ()
+        for v, count in self.variable_count:
+            if hasattr(v, "shape"):
+                for i in range(count):
+                    shape += v.shape
+        if hasattr(self.expr, "shape"):
+            shape += self.expr.shape
+        return shape
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @classmethod
+    def _get_zero_with_shape_like(cls, expr):
+        if isinstance(expr, (MatrixBase, NDimArray)):
+            return expr.zeros(*expr.shape)
+        elif isinstance(expr, MatrixExpr):
+            return ZeroMatrix(*expr.shape)
+        else:
+            raise RuntimeError("Unable to determine shape of array-derivative.")
+
+    @staticmethod
+    def _call_derive_scalar_by_matrix(expr: Expr, v: MatrixBase) -> Expr:
+        return v.applyfunc(lambda x: expr.diff(x))
+
+    @staticmethod
+    def _call_derive_scalar_by_matexpr(expr: Expr, v: MatrixExpr) -> Expr:
+        if expr.has(v):
+            return _matrix_derivative(expr, v)
+        else:
+            return ZeroMatrix(*v.shape)
+
+    @staticmethod
+    def _call_derive_scalar_by_array(expr: Expr, v: NDimArray) -> Expr:
+        return v.applyfunc(lambda x: expr.diff(x))
+
+    @staticmethod
+    def _call_derive_matrix_by_scalar(expr: MatrixBase, v: Expr) -> Expr:
+        return _matrix_derivative(expr, v)
+
+    @staticmethod
+    def _call_derive_matexpr_by_scalar(expr: MatrixExpr, v: Expr) -> Expr:
+        return expr._eval_derivative(v)
+
+    @staticmethod
+    def _call_derive_array_by_scalar(expr: NDimArray, v: Expr) -> Expr:
+        return expr.applyfunc(lambda x: x.diff(v))
+
+    @staticmethod
+    def _call_derive_default(expr: Expr, v: Expr) -> Expr | None:
+        if expr.has(v):
+            return _matrix_derivative(expr, v)
+        else:
+            return None
+
+    @classmethod
+    def _dispatch_eval_derivative_n_times(cls, expr, v, count):
+        # Evaluate the derivative `n` times.  If
+        # `_eval_derivative_n_times` is not overridden by the current
+        # object, the default in `Basic` will call a loop over
+        # `_eval_derivative`:
+
+        if not isinstance(count, (int, Integer)) or ((count <= 0) == True):
+            return None
+
+        # TODO: this could be done with multiple-dispatching:
+        if expr.is_scalar:
+            if isinstance(v, MatrixBase):
+                result = cls._call_derive_scalar_by_matrix(expr, v)
+            elif isinstance(v, MatrixExpr):
+                result = cls._call_derive_scalar_by_matexpr(expr, v)
+            elif isinstance(v, NDimArray):
+                result = cls._call_derive_scalar_by_array(expr, v)
+            elif v.is_scalar:
+                # scalar by scalar has a special
+                return super()._dispatch_eval_derivative_n_times(expr, v, count)
+            else:
+                return None
+        elif v.is_scalar:
+            if isinstance(expr, MatrixBase):
+                result = cls._call_derive_matrix_by_scalar(expr, v)
+            elif isinstance(expr, MatrixExpr):
+                result = cls._call_derive_matexpr_by_scalar(expr, v)
+            elif isinstance(expr, NDimArray):
+                result = cls._call_derive_array_by_scalar(expr, v)
+            else:
+                return None
+        else:
+            # Both `expr` and `v` are some array/matrix type:
+            if isinstance(expr, MatrixBase) or isinstance(v, MatrixBase):
+                result = derive_by_array(expr, v)
+            elif isinstance(expr, MatrixExpr) and isinstance(v, MatrixExpr):
+                result = cls._call_derive_default(expr, v)
+            elif isinstance(expr, MatrixExpr) or isinstance(v, MatrixExpr):
+                # if one expression is a symbolic matrix expression while the other isn't, don't evaluate:
+                return None
+            else:
+                result = derive_by_array(expr, v)
+        if result is None:
+            return None
+        if count == 1:
+            return result
+        else:
+            return cls._dispatch_eval_derivative_n_times(result, v, count - 1)

.venv/lib/python3.11/site-packages/sympy/tensor/array/arrayop.py

@@ -0,0 +1,528 @@
+import itertools
+from collections.abc import Iterable
+
+from sympy.core._print_helpers import Printable
+from sympy.core.containers import Tuple
+from sympy.core.function import diff
+from sympy.core.singleton import S
+from sympy.core.sympify import _sympify
+
+from sympy.tensor.array.ndim_array import NDimArray
+from sympy.tensor.array.dense_ndim_array import DenseNDimArray, ImmutableDenseNDimArray
+from sympy.tensor.array.sparse_ndim_array import SparseNDimArray
+
+
+def _arrayfy(a):
+    from sympy.matrices import MatrixBase
+
+    if isinstance(a, NDimArray):
+        return a
+    if isinstance(a, (MatrixBase, list, tuple, Tuple)):
+        return ImmutableDenseNDimArray(a)
+    return a
+
+
+def tensorproduct(*args):
+    """
+    Tensor product among scalars or array-like objects.
+
+    The equivalent operator for array expressions is ``ArrayTensorProduct``,
+    which can be used to keep the expression unevaluated.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array import tensorproduct, Array
+    >>> from sympy.abc import x, y, z, t
+    >>> A = Array([[1, 2], [3, 4]])
+    >>> B = Array([x, y])
+    >>> tensorproduct(A, B)
+    [[[x, y], [2*x, 2*y]], [[3*x, 3*y], [4*x, 4*y]]]
+    >>> tensorproduct(A, x)
+    [[x, 2*x], [3*x, 4*x]]
+    >>> tensorproduct(A, B, B)
+    [[[[x**2, x*y], [x*y, y**2]], [[2*x**2, 2*x*y], [2*x*y, 2*y**2]]], [[[3*x**2, 3*x*y], [3*x*y, 3*y**2]], [[4*x**2, 4*x*y], [4*x*y, 4*y**2]]]]
+
+    Applying this function on two matrices will result in a rank 4 array.
+
+    >>> from sympy import Matrix, eye
+    >>> m = Matrix([[x, y], [z, t]])
+    >>> p = tensorproduct(eye(3), m)
+    >>> p
+    [[[[x, y], [z, t]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[x, y], [z, t]], [[0, 0], [0, 0]]], [[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[x, y], [z, t]]]]
+
+    See Also
+    ========
+
+    sympy.tensor.array.expressions.array_expressions.ArrayTensorProduct
+
+    """
+    from sympy.tensor.array import SparseNDimArray, ImmutableSparseNDimArray
+
+    if len(args) == 0:
+        return S.One
+    if len(args) == 1:
+        return _arrayfy(args[0])
+    from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
+    from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct
+    from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
+    from sympy.matrices.expressions.matexpr import MatrixSymbol
+    if any(isinstance(arg, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)) for arg in args):
+        return ArrayTensorProduct(*args)
+    if len(args) > 2:
+        return tensorproduct(tensorproduct(args[0], args[1]), *args[2:])
+
+    # length of args is 2:
+    a, b = map(_arrayfy, args)
+
+    if not isinstance(a, NDimArray) or not isinstance(b, NDimArray):
+        return a*b
+
+    if isinstance(a, SparseNDimArray) and isinstance(b, SparseNDimArray):
+        lp = len(b)
+        new_array = {k1*lp + k2: v1*v2 for k1, v1 in a._sparse_array.items() for k2, v2 in b._sparse_array.items()}
+        return ImmutableSparseNDimArray(new_array, a.shape + b.shape)
+
+    product_list = [i*j for i in Flatten(a) for j in Flatten(b)]
+    return ImmutableDenseNDimArray(product_list, a.shape + b.shape)
+
+
+def _util_contraction_diagonal(array, *contraction_or_diagonal_axes):
+    array = _arrayfy(array)
+
+    # Verify contraction_axes:
+    taken_dims = set()
+    for axes_group in contraction_or_diagonal_axes:
+        if not isinstance(axes_group, Iterable):
+            raise ValueError("collections of contraction/diagonal axes expected")
+
+        dim = array.shape[axes_group[0]]
+
+        for d in axes_group:
+            if d in taken_dims:
+                raise ValueError("dimension specified more than once")
+            if dim != array.shape[d]:
+                raise ValueError("cannot contract or diagonalize between axes of different dimension")
+            taken_dims.add(d)
+
+    rank = array.rank()
+
+    remaining_shape = [dim for i, dim in enumerate(array.shape) if i not in taken_dims]
+    cum_shape = [0]*rank
+    _cumul = 1
+    for i in range(rank):
+        cum_shape[rank - i - 1] = _cumul
+        _cumul *= int(array.shape[rank - i - 1])
+
+    # DEFINITION: by absolute position it is meant the position along the one
+    # dimensional array containing all the tensor components.
+
+    # Possible future work on this module: move computation of absolute
+    # positions to a class method.
+
+    # Determine absolute positions of the uncontracted indices:
+    remaining_indices = [[cum_shape[i]*j for j in range(array.shape[i])]
+                         for i in range(rank) if i not in taken_dims]
+
+    # Determine absolute positions of the contracted indices:
+    summed_deltas = []
+    for axes_group in contraction_or_diagonal_axes:
+        lidx = []
+        for js in range(array.shape[axes_group[0]]):
+            lidx.append(sum(cum_shape[ig] * js for ig in axes_group))
+        summed_deltas.append(lidx)
+
+    return array, remaining_indices, remaining_shape, summed_deltas
+
+
+def tensorcontraction(array, *contraction_axes):
+    """
+    Contraction of an array-like object on the specified axes.
+
+    The equivalent operator for array expressions is ``ArrayContraction``,
+    which can be used to keep the expression unevaluated.
+
+    Examples
+    ========
+
+    >>> from sympy import Array, tensorcontraction
+    >>> from sympy import Matrix, eye
+    >>> tensorcontraction(eye(3), (0, 1))
+    3
+    >>> A = Array(range(18), (3, 2, 3))
+    >>> A
+    [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
+    >>> tensorcontraction(A, (0, 2))
+    [21, 30]
+
+    Matrix multiplication may be emulated with a proper combination of
+    ``tensorcontraction`` and ``tensorproduct``
+
+    >>> from sympy import tensorproduct
+    >>> from sympy.abc import a,b,c,d,e,f,g,h
+    >>> m1 = Matrix([[a, b], [c, d]])
+    >>> m2 = Matrix([[e, f], [g, h]])
+    >>> p = tensorproduct(m1, m2)
+    >>> p
+    [[[[a*e, a*f], [a*g, a*h]], [[b*e, b*f], [b*g, b*h]]], [[[c*e, c*f], [c*g, c*h]], [[d*e, d*f], [d*g, d*h]]]]
+    >>> tensorcontraction(p, (1, 2))
+    [[a*e + b*g, a*f + b*h], [c*e + d*g, c*f + d*h]]
+    >>> m1*m2
+    Matrix([
+    [a*e + b*g, a*f + b*h],
+    [c*e + d*g, c*f + d*h]])
+
+    See Also
+    ========
+
+    sympy.tensor.array.expressions.array_expressions.ArrayContraction
+
+    """
+    from sympy.tensor.array.expressions.array_expressions import _array_contraction
+    from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
+    from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
+    from sympy.matrices.expressions.matexpr import MatrixSymbol
+    if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
+        return _array_contraction(array, *contraction_axes)
+
+    array, remaining_indices, remaining_shape, summed_deltas = _util_contraction_diagonal(array, *contraction_axes)
+
+    # Compute the contracted array:
+    #
+    # 1. external for loops on all uncontracted indices.
+    #    Uncontracted indices are determined by the combinatorial product of
+    #    the absolute positions of the remaining indices.
+    # 2. internal loop on all contracted indices.
+    #    It sums the values of the absolute contracted index and the absolute
+    #    uncontracted index for the external loop.
+    contracted_array = []
+    for icontrib in itertools.product(*remaining_indices):
+        index_base_position = sum(icontrib)
+        isum = S.Zero
+        for sum_to_index in itertools.product(*summed_deltas):
+            idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
+            isum += array[idx]
+
+        contracted_array.append(isum)
+
+    if len(remaining_indices) == 0:
+        assert len(contracted_array) == 1
+        return contracted_array[0]
+
+    return type(array)(contracted_array, remaining_shape)
+
+
+def tensordiagonal(array, *diagonal_axes):
+    """
+    Diagonalization of an array-like object on the specified axes.
+
+    This is equivalent to multiplying the expression by Kronecker deltas
+    uniting the axes.
+
+    The diagonal indices are put at the end of the axes.
+
+    The equivalent operator for array expressions is ``ArrayDiagonal``, which
+    can be used to keep the expression unevaluated.
+
+    Examples
+    ========
+
+    ``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is
+    equivalent to the diagonal of the matrix:
+
+    >>> from sympy import Array, tensordiagonal
+    >>> from sympy import Matrix, eye
+    >>> tensordiagonal(eye(3), (0, 1))
+    [1, 1, 1]
+
+    >>> from sympy.abc import a,b,c,d
+    >>> m1 = Matrix([[a, b], [c, d]])
+    >>> tensordiagonal(m1, [0, 1])
+    [a, d]
+
+    In case of higher dimensional arrays, the diagonalized out dimensions
+    are appended removed and appended as a single dimension at the end:
+
+    >>> A = Array(range(18), (3, 2, 3))
+    >>> A
+    [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
+    >>> tensordiagonal(A, (0, 2))
+    [[0, 7, 14], [3, 10, 17]]
+    >>> from sympy import permutedims
+    >>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0])
+    True
+
+    See Also
+    ========
+
+    sympy.tensor.array.expressions.array_expressions.ArrayDiagonal
+
+    """
+    if any(len(i) <= 1 for i in diagonal_axes):
+        raise ValueError("need at least two axes to diagonalize")
+
+    from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
+    from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
+    from sympy.tensor.array.expressions.array_expressions import ArrayDiagonal, _array_diagonal
+    from sympy.matrices.expressions.matexpr import MatrixSymbol
+    if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
+        return _array_diagonal(array, *diagonal_axes)
+
+    ArrayDiagonal._validate(array, *diagonal_axes)
+
+    array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(array, *diagonal_axes)
+
+    # Compute the diagonalized array:
+    #
+    # 1. external for loops on all undiagonalized indices.
+    #    Undiagonalized indices are determined by the combinatorial product of
+    #    the absolute positions of the remaining indices.
+    # 2. internal loop on all diagonal indices.
+    #    It appends the values of the absolute diagonalized index and the absolute
+    #    undiagonalized index for the external loop.
+    diagonalized_array = []
+    diagonal_shape = [len(i) for i in diagonal_deltas]
+    for icontrib in itertools.product(*remaining_indices):
+        index_base_position = sum(icontrib)
+        isum = []
+        for sum_to_index in itertools.product(*diagonal_deltas):
+            idx = array._get_tuple_index(index_base_position + sum(sum_to_index))
+            isum.append(array[idx])
+
+        isum = type(array)(isum).reshape(*diagonal_shape)
+        diagonalized_array.append(isum)
+
+    return type(array)(diagonalized_array, remaining_shape + diagonal_shape)
+
+
+def derive_by_array(expr, dx):
+    r"""
+    Derivative by arrays. Supports both arrays and scalars.
+
+    The equivalent operator for array expressions is ``array_derive``.
+
+    Explanation
+    ===========
+
+    Given the array `A_{i_1, \ldots, i_N}` and the array `X_{j_1, \ldots, j_M}`
+    this function will return a new array `B` defined by
+
+    `B_{j_1,\ldots,j_M,i_1,\ldots,i_N} := \frac{\partial A_{i_1,\ldots,i_N}}{\partial X_{j_1,\ldots,j_M}}`
+
+    Examples
+    ========
+
+    >>> from sympy import derive_by_array
+    >>> from sympy.abc import x, y, z, t
+    >>> from sympy import cos
+    >>> derive_by_array(cos(x*t), x)
+    -t*sin(t*x)
+    >>> derive_by_array(cos(x*t), [x, y, z, t])
+    [-t*sin(t*x), 0, 0, -x*sin(t*x)]
+    >>> derive_by_array([x, y**2*z], [[x, y], [z, t]])
+    [[[1, 0], [0, 2*y*z]], [[0, y**2], [0, 0]]]
+
+    """
+    from sympy.matrices import MatrixBase
+    from sympy.tensor.array import SparseNDimArray
+    array_types = (Iterable, MatrixBase, NDimArray)
+
+    if isinstance(dx, array_types):
+        dx = ImmutableDenseNDimArray(dx)
+        for i in dx:
+            if not i._diff_wrt:
+                raise ValueError("cannot derive by this array")
+
+    if isinstance(expr, array_types):
+        if isinstance(expr, NDimArray):
+            expr = expr.as_immutable()
+        else:
+            expr = ImmutableDenseNDimArray(expr)
+
+        if isinstance(dx, array_types):
+            if isinstance(expr, SparseNDimArray):
+                lp = len(expr)
+                new_array = {k + i*lp: v
+                             for i, x in enumerate(Flatten(dx))
+                             for k, v in expr.diff(x)._sparse_array.items()}
+            else:
+                new_array = [[y.diff(x) for y in Flatten(expr)] for x in Flatten(dx)]
+            return type(expr)(new_array, dx.shape + expr.shape)
+        else:
+            return expr.diff(dx)
+    else:
+        expr = _sympify(expr)
+        if isinstance(dx, array_types):
+            return ImmutableDenseNDimArray([expr.diff(i) for i in Flatten(dx)], dx.shape)
+        else:
+            dx = _sympify(dx)
+            return diff(expr, dx)
+
+
+def permutedims(expr, perm=None, index_order_old=None, index_order_new=None):
+    """
+    Permutes the indices of an array.
+
+    Parameter specifies the permutation of the indices.
+
+    The equivalent operator for array expressions is ``PermuteDims``, which can
+    be used to keep the expression unevaluated.
+
+    Examples
+    ========
+
+    >>> from sympy.abc import x, y, z, t
+    >>> from sympy import sin
+    >>> from sympy import Array, permutedims
+    >>> a = Array([[x, y, z], [t, sin(x), 0]])
+    >>> a
+    [[x, y, z], [t, sin(x), 0]]
+    >>> permutedims(a, (1, 0))
+    [[x, t], [y, sin(x)], [z, 0]]
+
+    If the array is of second order, ``transpose`` can be used:
+
+    >>> from sympy import transpose
+    >>> transpose(a)
+    [[x, t], [y, sin(x)], [z, 0]]
+
+    Examples on higher dimensions:
+
+    >>> b = Array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    >>> permutedims(b, (2, 1, 0))
+    [[[1, 5], [3, 7]], [[2, 6], [4, 8]]]
+    >>> permutedims(b, (1, 2, 0))
+    [[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
+
+    An alternative way to specify the same permutations as in the previous
+    lines involves passing the *old* and *new* indices, either as a list or as
+    a string:
+
+    >>> permutedims(b, index_order_old="cba", index_order_new="abc")
+    [[[1, 5], [3, 7]], [[2, 6], [4, 8]]]
+    >>> permutedims(b, index_order_old="cab", index_order_new="abc")
+    [[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
+
+    ``Permutation`` objects are also allowed:
+
+    >>> from sympy.combinatorics import Permutation
+    >>> permutedims(b, Permutation([1, 2, 0]))
+    [[[1, 5], [2, 6]], [[3, 7], [4, 8]]]
+
+    See Also
+    ========
+
+    sympy.tensor.array.expressions.array_expressions.PermuteDims
+
+    """
+    from sympy.tensor.array import SparseNDimArray
+
+    from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
+    from sympy.tensor.array.expressions.array_expressions import _CodegenArrayAbstract
+    from sympy.tensor.array.expressions.array_expressions import _permute_dims
+    from sympy.matrices.expressions.matexpr import MatrixSymbol
+    from sympy.tensor.array.expressions import PermuteDims
+    from sympy.tensor.array.expressions.array_expressions import get_rank
+    perm = PermuteDims._get_permutation_from_arguments(perm, index_order_old, index_order_new, get_rank(expr))
+    if isinstance(expr, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
+        return _permute_dims(expr, perm)
+
+    if not isinstance(expr, NDimArray):
+        expr = ImmutableDenseNDimArray(expr)
+
+    from sympy.combinatorics import Permutation
+    if not isinstance(perm, Permutation):
+        perm = Permutation(list(perm))
+
+    if perm.size != expr.rank():
+        raise ValueError("wrong permutation size")
+
+    # Get the inverse permutation:
+    iperm = ~perm
+    new_shape = perm(expr.shape)
+
+    if isinstance(expr, SparseNDimArray):
+        return type(expr)({tuple(perm(expr._get_tuple_index(k))): v
+                           for k, v in expr._sparse_array.items()}, new_shape)
+
+    indices_span = perm([range(i) for i in expr.shape])
+
+    new_array = [None]*len(expr)
+    for i, idx in enumerate(itertools.product(*indices_span)):
+        t = iperm(idx)
+        new_array[i] = expr[t]
+
+    return type(expr)(new_array, new_shape)
+
+
+class Flatten(Printable):
+    """
+    Flatten an iterable object to a list in a lazy-evaluation way.
+
+    Notes
+    =====
+
+    This class is an iterator with which the memory cost can be economised.
+    Optimisation has been considered to ameliorate the performance for some
+    specific data types like DenseNDimArray and SparseNDimArray.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array.arrayop import Flatten
+    >>> from sympy.tensor.array import Array
+    >>> A = Array(range(6)).reshape(2, 3)
+    >>> Flatten(A)
+    Flatten([[0, 1, 2], [3, 4, 5]])
+    >>> [i for i in Flatten(A)]
+    [0, 1, 2, 3, 4, 5]
+    """
+    def __init__(self, iterable):
+        from sympy.matrices.matrixbase import MatrixBase
+        from sympy.tensor.array import NDimArray
+
+        if not isinstance(iterable, (Iterable, MatrixBase)):
+            raise NotImplementedError("Data type not yet supported")
+
+        if isinstance(iterable, list):
+            iterable = NDimArray(iterable)
+
+        self._iter = iterable
+        self._idx = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        from sympy.matrices.matrixbase import MatrixBase
+
+        if len(self._iter) > self._idx:
+            if isinstance(self._iter, DenseNDimArray):
+                result = self._iter._array[self._idx]
+
+            elif isinstance(self._iter, SparseNDimArray):
+                if self._idx in self._iter._sparse_array:
+                    result = self._iter._sparse_array[self._idx]
+                else:
+                    result = 0
+
+            elif isinstance(self._iter, MatrixBase):
+                result = self._iter[self._idx]
+
+            elif hasattr(self._iter, '__next__'):
+                result = next(self._iter)
+
+            else:
+                result = self._iter[self._idx]
+
+        else:
+            raise StopIteration
+
+        self._idx += 1
+        return result
+
+    def next(self):
+        return self.__next__()
+
+    def _sympystr(self, printer):
+        return type(self).__name__ + '(' + printer._print(self._iter) + ')'

.venv/lib/python3.11/site-packages/sympy/tensor/array/dense_ndim_array.py

@@ -0,0 +1,206 @@
+import functools
+from typing import List
+
+from sympy.core.basic import Basic
+from sympy.core.containers import Tuple
+from sympy.core.singleton import S
+from sympy.core.sympify import _sympify
+from sympy.tensor.array.mutable_ndim_array import MutableNDimArray
+from sympy.tensor.array.ndim_array import NDimArray, ImmutableNDimArray, ArrayKind
+from sympy.utilities.iterables import flatten
+
+
+class DenseNDimArray(NDimArray):
+
+    _array: List[Basic]
+
+    def __new__(self, *args, **kwargs):
+        return ImmutableDenseNDimArray(*args, **kwargs)
+
+    @property
+    def kind(self) -> ArrayKind:
+        return ArrayKind._union(self._array)
+
+    def __getitem__(self, index):
+        """
+        Allows to get items from N-dim array.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray([0, 1, 2, 3], (2, 2))
+        >>> a
+        [[0, 1], [2, 3]]
+        >>> a[0, 0]
+        0
+        >>> a[1, 1]
+        3
+        >>> a[0]
+        [0, 1]
+        >>> a[1]
+        [2, 3]
+
+
+        Symbolic index:
+
+        >>> from sympy.abc import i, j
+        >>> a[i, j]
+        [[0, 1], [2, 3]][i, j]
+
+        Replace `i` and `j` to get element `(1, 1)`:
+
+        >>> a[i, j].subs({i: 1, j: 1})
+        3
+
+        """
+        syindex = self._check_symbolic_index(index)
+        if syindex is not None:
+            return syindex
+
+        index = self._check_index_for_getitem(index)
+
+        if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
+            sl_factors, eindices = self._get_slice_data_for_array_access(index)
+            array = [self._array[self._parse_index(i)] for i in eindices]
+            nshape = [len(el) for i, el in enumerate(sl_factors) if isinstance(index[i], slice)]
+            return type(self)(array, nshape)
+        else:
+            index = self._parse_index(index)
+            return self._array[index]
+
+    @classmethod
+    def zeros(cls, *shape):
+        list_length = functools.reduce(lambda x, y: x*y, shape, S.One)
+        return cls._new(([0]*list_length,), shape)
+
+    def tomatrix(self):
+        """
+        Converts MutableDenseNDimArray to Matrix. Can convert only 2-dim array, else will raise error.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray([1 for i in range(9)], (3, 3))
+        >>> b = a.tomatrix()
+        >>> b
+        Matrix([
+        [1, 1, 1],
+        [1, 1, 1],
+        [1, 1, 1]])
+
+        """
+        from sympy.matrices import Matrix
+
+        if self.rank() != 2:
+            raise ValueError('Dimensions must be of size of 2')
+
+        return Matrix(self.shape[0], self.shape[1], self._array)
+
+    def reshape(self, *newshape):
+        """
+        Returns MutableDenseNDimArray instance with new shape. Elements number
+        must be        suitable to new shape. The only argument of method sets
+        new shape.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3))
+        >>> a.shape
+        (2, 3)
+        >>> a
+        [[1, 2, 3], [4, 5, 6]]
+        >>> b = a.reshape(3, 2)
+        >>> b.shape
+        (3, 2)
+        >>> b
+        [[1, 2], [3, 4], [5, 6]]
+
+        """
+        new_total_size = functools.reduce(lambda x,y: x*y, newshape)
+        if new_total_size != self._loop_size:
+            raise ValueError('Expecting reshape size to %d but got prod(%s) = %d' % (
+                self._loop_size, str(newshape), new_total_size))
+
+        # there is no `.func` as this class does not subtype `Basic`:
+        return type(self)(self._array, newshape)
+
+
+class ImmutableDenseNDimArray(DenseNDimArray, ImmutableNDimArray): # type: ignore
+    def __new__(cls, iterable, shape=None, **kwargs):
+        return cls._new(iterable, shape, **kwargs)
+
+    @classmethod
+    def _new(cls, iterable, shape, **kwargs):
+        shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
+        shape = Tuple(*map(_sympify, shape))
+        cls._check_special_bounds(flat_list, shape)
+        flat_list = flatten(flat_list)
+        flat_list = Tuple(*flat_list)
+        self = Basic.__new__(cls, flat_list, shape, **kwargs)
+        self._shape = shape
+        self._array = list(flat_list)
+        self._rank = len(shape)
+        self._loop_size = functools.reduce(lambda x,y: x*y, shape, 1)
+        return self
+
+    def __setitem__(self, index, value):
+        raise TypeError('immutable N-dim array')
+
+    def as_mutable(self):
+        return MutableDenseNDimArray(self)
+
+    def _eval_simplify(self, **kwargs):
+        from sympy.simplify.simplify import simplify
+        return self.applyfunc(simplify)
+
+class MutableDenseNDimArray(DenseNDimArray, MutableNDimArray):
+
+    def __new__(cls, iterable=None, shape=None, **kwargs):
+        return cls._new(iterable, shape, **kwargs)
+
+    @classmethod
+    def _new(cls, iterable, shape, **kwargs):
+        shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
+        flat_list = flatten(flat_list)
+        self = object.__new__(cls)
+        self._shape = shape
+        self._array = list(flat_list)
+        self._rank = len(shape)
+        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
+        return self
+
+    def __setitem__(self, index, value):
+        """Allows to set items to MutableDenseNDimArray.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray.zeros(2,  2)
+        >>> a[0,0] = 1
+        >>> a[1,1] = 1
+        >>> a
+        [[1, 0], [0, 1]]
+
+        """
+        if isinstance(index, tuple) and any(isinstance(i, slice) for i in index):
+            value, eindices, slice_offsets = self._get_slice_data_for_array_assignment(index, value)
+            for i in eindices:
+                other_i = [ind - j for ind, j in zip(i, slice_offsets) if j is not None]
+                self._array[self._parse_index(i)] = value[other_i]
+        else:
+            index = self._parse_index(index)
+            self._setter_iterable_check(value)
+            value = _sympify(value)
+            self._array[index] = value
+
+    def as_immutable(self):
+        return ImmutableDenseNDimArray(self)
+
+    @property
+    def free_symbols(self):
+        return {i for j in self._array for i in j.free_symbols}

.venv/lib/python3.11/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

.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/array_expressions.py

@@ -0,0 +1,1967 @@
+import collections.abc
+import operator
+from collections import defaultdict, Counter
+from functools import reduce
+import itertools
+from itertools import accumulate
+from typing import Optional, List, Tuple as tTuple
+
+import typing
+
+from sympy.core.numbers import Integer
+from sympy.core.relational import Equality
+from sympy.functions.special.tensor_functions import KroneckerDelta
+from sympy.core.basic import Basic
+from sympy.core.containers import Tuple
+from sympy.core.expr import Expr
+from sympy.core.function import (Function, Lambda)
+from sympy.core.mul import Mul
+from sympy.core.singleton import S
+from sympy.core.sorting import default_sort_key
+from sympy.core.symbol import (Dummy, Symbol)
+from sympy.matrices.matrixbase import MatrixBase
+from sympy.matrices.expressions.diagonal import diagonalize_vector
+from sympy.matrices.expressions.matexpr import MatrixExpr
+from sympy.matrices.expressions.special import ZeroMatrix
+from sympy.tensor.array.arrayop import (permutedims, tensorcontraction, tensordiagonal, tensorproduct)
+from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
+from sympy.tensor.array.ndim_array import NDimArray
+from sympy.tensor.indexed import (Indexed, IndexedBase)
+from sympy.matrices.expressions.matexpr import MatrixElement
+from sympy.tensor.array.expressions.utils import _apply_recursively_over_nested_lists, _sort_contraction_indices, \
+    _get_mapping_from_subranks, _build_push_indices_up_func_transformation, _get_contraction_links, \
+    _build_push_indices_down_func_transformation
+from sympy.combinatorics import Permutation
+from sympy.combinatorics.permutations import _af_invert
+from sympy.core.sympify import _sympify
+
+
+class _ArrayExpr(Expr):
+    shape: tTuple[Expr, ...]
+
+    def __getitem__(self, item):
+        if not isinstance(item, collections.abc.Iterable):
+            item = (item,)
+        ArrayElement._check_shape(self, item)
+        return self._get(item)
+
+    def _get(self, item):
+        return _get_array_element_or_slice(self, item)
+
+
+class ArraySymbol(_ArrayExpr):
+    """
+    Symbol representing an array expression
+    """
+
+    def __new__(cls, symbol, shape: typing.Iterable) -> "ArraySymbol":
+        if isinstance(symbol, str):
+            symbol = Symbol(symbol)
+        # symbol = _sympify(symbol)
+        shape = Tuple(*map(_sympify, shape))
+        obj = Expr.__new__(cls, symbol, shape)
+        return obj
+
+    @property
+    def name(self):
+        return self._args[0]
+
+    @property
+    def shape(self):
+        return self._args[1]
+
+    def as_explicit(self):
+        if not all(i.is_Integer for i in self.shape):
+            raise ValueError("cannot express explicit array with symbolic shape")
+        data = [self[i] for i in itertools.product(*[range(j) for j in self.shape])]
+        return ImmutableDenseNDimArray(data).reshape(*self.shape)
+
+
+class ArrayElement(Expr):
+    """
+    An element of an array.
+    """
+
+    _diff_wrt = True
+    is_symbol = True
+    is_commutative = True
+
+    def __new__(cls, name, indices):
+        if isinstance(name, str):
+            name = Symbol(name)
+        name = _sympify(name)
+        if not isinstance(indices, collections.abc.Iterable):
+            indices = (indices,)
+        indices = _sympify(tuple(indices))
+        cls._check_shape(name, indices)
+        obj = Expr.__new__(cls, name, indices)
+        return obj
+
+    @classmethod
+    def _check_shape(cls, name, indices):
+        indices = tuple(indices)
+        if hasattr(name, "shape"):
+            index_error = IndexError("number of indices does not match shape of the array")
+            if len(indices) != len(name.shape):
+                raise index_error
+            if any((i >= s) == True for i, s in zip(indices, name.shape)):
+                raise ValueError("shape is out of bounds")
+        if any((i < 0) == True for i in indices):
+            raise ValueError("shape contains negative values")
+
+    @property
+    def name(self):
+        return self._args[0]
+
+    @property
+    def indices(self):
+        return self._args[1]
+
+    def _eval_derivative(self, s):
+        if not isinstance(s, ArrayElement):
+            return S.Zero
+
+        if s == self:
+            return S.One
+
+        if s.name != self.name:
+            return S.Zero
+
+        return Mul.fromiter(KroneckerDelta(i, j) for i, j in zip(self.indices, s.indices))
+
+
+class ZeroArray(_ArrayExpr):
+    """
+    Symbolic array of zeros. Equivalent to ``ZeroMatrix`` for matrices.
+    """
+
+    def __new__(cls, *shape):
+        if len(shape) == 0:
+            return S.Zero
+        shape = map(_sympify, shape)
+        obj = Expr.__new__(cls, *shape)
+        return obj
+
+    @property
+    def shape(self):
+        return self._args
+
+    def as_explicit(self):
+        if not all(i.is_Integer for i in self.shape):
+            raise ValueError("Cannot return explicit form for symbolic shape.")
+        return ImmutableDenseNDimArray.zeros(*self.shape)
+
+    def _get(self, item):
+        return S.Zero
+
+
+class OneArray(_ArrayExpr):
+    """
+    Symbolic array of ones.
+    """
+
+    def __new__(cls, *shape):
+        if len(shape) == 0:
+            return S.One
+        shape = map(_sympify, shape)
+        obj = Expr.__new__(cls, *shape)
+        return obj
+
+    @property
+    def shape(self):
+        return self._args
+
+    def as_explicit(self):
+        if not all(i.is_Integer for i in self.shape):
+            raise ValueError("Cannot return explicit form for symbolic shape.")
+        return ImmutableDenseNDimArray([S.One for i in range(reduce(operator.mul, self.shape))]).reshape(*self.shape)
+
+    def _get(self, item):
+        return S.One
+
+
+class _CodegenArrayAbstract(Basic):
+
+    @property
+    def subranks(self):
+        """
+        Returns the ranks of the objects in the uppermost tensor product inside
+        the current object.  In case no tensor products are contained, return
+        the atomic ranks.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import tensorproduct, tensorcontraction
+        >>> from sympy import MatrixSymbol
+        >>> M = MatrixSymbol("M", 3, 3)
+        >>> N = MatrixSymbol("N", 3, 3)
+        >>> P = MatrixSymbol("P", 3, 3)
+
+        Important: do not confuse the rank of the matrix with the rank of an array.
+
+        >>> tp = tensorproduct(M, N, P)
+        >>> tp.subranks
+        [2, 2, 2]
+
+        >>> co = tensorcontraction(tp, (1, 2), (3, 4))
+        >>> co.subranks
+        [2, 2, 2]
+        """
+        return self._subranks[:]
+
+    def subrank(self):
+        """
+        The sum of ``subranks``.
+        """
+        return sum(self.subranks)
+
+    @property
+    def shape(self):
+        return self._shape
+
+    def doit(self, **hints):
+        deep = hints.get("deep", True)
+        if deep:
+            return self.func(*[arg.doit(**hints) for arg in self.args])._canonicalize()
+        else:
+            return self._canonicalize()
+
+class ArrayTensorProduct(_CodegenArrayAbstract):
+    r"""
+    Class to represent the tensor product of array-like objects.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        args = [_sympify(arg) for arg in args]
+
+        canonicalize = kwargs.pop("canonicalize", False)
+
+        ranks = [get_rank(arg) for arg in args]
+
+        obj = Basic.__new__(cls, *args)
+        obj._subranks = ranks
+        shapes = [get_shape(i) for i in args]
+
+        if any(i is None for i in shapes):
+            obj._shape = None
+        else:
+            obj._shape = tuple(j for i in shapes for j in i)
+        if canonicalize:
+            return obj._canonicalize()
+        return obj
+
+    def _canonicalize(self):
+        args = self.args
+        args = self._flatten(args)
+
+        ranks = [get_rank(arg) for arg in args]
+
+        # Check if there are nested permutation and lift them up:
+        permutation_cycles = []
+        for i, arg in enumerate(args):
+            if not isinstance(arg, PermuteDims):
+                continue
+            permutation_cycles.extend([[k + sum(ranks[:i]) for k in j] for j in arg.permutation.cyclic_form])
+            args[i] = arg.expr
+        if permutation_cycles:
+            return _permute_dims(_array_tensor_product(*args), Permutation(sum(ranks)-1)*Permutation(permutation_cycles))
+
+        if len(args) == 1:
+            return args[0]
+
+        # If any object is a ZeroArray, return a ZeroArray:
+        if any(isinstance(arg, (ZeroArray, ZeroMatrix)) for arg in args):
+            shapes = reduce(operator.add, [get_shape(i) for i in args], ())
+            return ZeroArray(*shapes)
+
+        # If there are contraction objects inside, transform the whole
+        # expression into `ArrayContraction`:
+        contractions = {i: arg for i, arg in enumerate(args) if isinstance(arg, ArrayContraction)}
+        if contractions:
+            ranks = [_get_subrank(arg) if isinstance(arg, ArrayContraction) else get_rank(arg) for arg in args]
+            cumulative_ranks = list(accumulate([0] + ranks))[:-1]
+            tp = _array_tensor_product(*[arg.expr if isinstance(arg, ArrayContraction) else arg for arg in args])
+            contraction_indices = [tuple(cumulative_ranks[i] + k for k in j) for i, arg in contractions.items() for j in arg.contraction_indices]
+            return _array_contraction(tp, *contraction_indices)
+
+        diagonals = {i: arg for i, arg in enumerate(args) if isinstance(arg, ArrayDiagonal)}
+        if diagonals:
+            inverse_permutation = []
+            last_perm = []
+            ranks = [get_rank(arg) for arg in args]
+            cumulative_ranks = list(accumulate([0] + ranks))[:-1]
+            for i, arg in enumerate(args):
+                if isinstance(arg, ArrayDiagonal):
+                    i1 = get_rank(arg) - len(arg.diagonal_indices)
+                    i2 = len(arg.diagonal_indices)
+                    inverse_permutation.extend([cumulative_ranks[i] + j for j in range(i1)])
+                    last_perm.extend([cumulative_ranks[i] + j for j in range(i1, i1 + i2)])
+                else:
+                    inverse_permutation.extend([cumulative_ranks[i] + j for j in range(get_rank(arg))])
+            inverse_permutation.extend(last_perm)
+            tp = _array_tensor_product(*[arg.expr if isinstance(arg, ArrayDiagonal) else arg for arg in args])
+            ranks2 = [_get_subrank(arg) if isinstance(arg, ArrayDiagonal) else get_rank(arg) for arg in args]
+            cumulative_ranks2 = list(accumulate([0] + ranks2))[:-1]
+            diagonal_indices = [tuple(cumulative_ranks2[i] + k for k in j) for i, arg in diagonals.items() for j in arg.diagonal_indices]
+            return _permute_dims(_array_diagonal(tp, *diagonal_indices), _af_invert(inverse_permutation))
+
+        return self.func(*args, canonicalize=False)
+
+    @classmethod
+    def _flatten(cls, args):
+        args = [i for arg in args for i in (arg.args if isinstance(arg, cls) else [arg])]
+        return args
+
+    def as_explicit(self):
+        return tensorproduct(*[arg.as_explicit() if hasattr(arg, "as_explicit") else arg for arg in self.args])
+
+
+class ArrayAdd(_CodegenArrayAbstract):
+    r"""
+    Class for elementwise array additions.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        args = [_sympify(arg) for arg in args]
+        ranks = [get_rank(arg) for arg in args]
+        ranks = list(set(ranks))
+        if len(ranks) != 1:
+            raise ValueError("summing arrays of different ranks")
+        shapes = [arg.shape for arg in args]
+        if len({i for i in shapes if i is not None}) > 1:
+            raise ValueError("mismatching shapes in addition")
+
+        canonicalize = kwargs.pop("canonicalize", False)
+
+        obj = Basic.__new__(cls, *args)
+        obj._subranks = ranks
+        if any(i is None for i in shapes):
+            obj._shape = None
+        else:
+            obj._shape = shapes[0]
+        if canonicalize:
+            return obj._canonicalize()
+        return obj
+
+    def _canonicalize(self):
+        args = self.args
+
+        # Flatten:
+        args = self._flatten_args(args)
+
+        shapes = [get_shape(arg) for arg in args]
+        args = [arg for arg in args if not isinstance(arg, (ZeroArray, ZeroMatrix))]
+        if len(args) == 0:
+            if any(i for i in shapes if i is None):
+                raise NotImplementedError("cannot handle addition of ZeroMatrix/ZeroArray and undefined shape object")
+            return ZeroArray(*shapes[0])
+        elif len(args) == 1:
+            return args[0]
+        return self.func(*args, canonicalize=False)
+
+    @classmethod
+    def _flatten_args(cls, args):
+        new_args = []
+        for arg in args:
+            if isinstance(arg, ArrayAdd):
+                new_args.extend(arg.args)
+            else:
+                new_args.append(arg)
+        return new_args
+
+    def as_explicit(self):
+        return reduce(
+            operator.add,
+            [arg.as_explicit() if hasattr(arg, "as_explicit") else arg for arg in self.args])
+
+
+class PermuteDims(_CodegenArrayAbstract):
+    r"""
+    Class to represent permutation of axes of arrays.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array import permutedims
+    >>> from sympy import MatrixSymbol
+    >>> M = MatrixSymbol("M", 3, 3)
+    >>> cg = permutedims(M, [1, 0])
+
+    The object ``cg`` represents the transposition of ``M``, as the permutation
+    ``[1, 0]`` will act on its indices by switching them:
+
+    `M_{ij} \Rightarrow M_{ji}`
+
+    This is evident when transforming back to matrix form:
+
+    >>> from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+    >>> convert_array_to_matrix(cg)
+    M.T
+
+    >>> N = MatrixSymbol("N", 3, 2)
+    >>> cg = permutedims(N, [1, 0])
+    >>> cg.shape
+    (2, 3)
+
+    There are optional parameters that can be used as alternative to the permutation:
+
+    >>> from sympy.tensor.array.expressions import ArraySymbol, PermuteDims
+    >>> M = ArraySymbol("M", (1, 2, 3, 4, 5))
+    >>> expr = PermuteDims(M, index_order_old="ijklm", index_order_new="kijml")
+    >>> expr
+    PermuteDims(M, (0 2 1)(3 4))
+    >>> expr.shape
+    (3, 1, 2, 5, 4)
+
+    Permutations of tensor products are simplified in order to achieve a
+    standard form:
+
+    >>> from sympy.tensor.array import tensorproduct
+    >>> M = MatrixSymbol("M", 4, 5)
+    >>> tp = tensorproduct(M, N)
+    >>> tp.shape
+    (4, 5, 3, 2)
+    >>> perm1 = permutedims(tp, [2, 3, 1, 0])
+
+    The args ``(M, N)`` have been sorted and the permutation has been
+    simplified, the expression is equivalent:
+
+    >>> perm1.expr.args
+    (N, M)
+    >>> perm1.shape
+    (3, 2, 5, 4)
+    >>> perm1.permutation
+    (2 3)
+
+    The permutation in its array form has been simplified from
+    ``[2, 3, 1, 0]`` to ``[0, 1, 3, 2]``, as the arguments of the tensor
+    product `M` and `N` have been switched:
+
+    >>> perm1.permutation.array_form
+    [0, 1, 3, 2]
+
+    We can nest a second permutation:
+
+    >>> perm2 = permutedims(perm1, [1, 0, 2, 3])
+    >>> perm2.shape
+    (2, 3, 5, 4)
+    >>> perm2.permutation.array_form
+    [1, 0, 3, 2]
+    """
+
+    def __new__(cls, expr, permutation=None, index_order_old=None, index_order_new=None, **kwargs):
+        from sympy.combinatorics import Permutation
+        expr = _sympify(expr)
+        expr_rank = get_rank(expr)
+        permutation = cls._get_permutation_from_arguments(permutation, index_order_old, index_order_new, expr_rank)
+        permutation = Permutation(permutation)
+        permutation_size = permutation.size
+        if permutation_size != expr_rank:
+            raise ValueError("Permutation size must be the length of the shape of expr")
+
+        canonicalize = kwargs.pop("canonicalize", False)
+
+        obj = Basic.__new__(cls, expr, permutation)
+        obj._subranks = [get_rank(expr)]
+        shape = get_shape(expr)
+        if shape is None:
+            obj._shape = None
+        else:
+            obj._shape = tuple(shape[permutation(i)] for i in range(len(shape)))
+        if canonicalize:
+            return obj._canonicalize()
+        return obj
+
+    def _canonicalize(self):
+        expr = self.expr
+        permutation = self.permutation
+        if isinstance(expr, PermuteDims):
+            subexpr = expr.expr
+            subperm = expr.permutation
+            permutation = permutation * subperm
+            expr = subexpr
+        if isinstance(expr, ArrayContraction):
+            expr, permutation = self._PermuteDims_denestarg_ArrayContraction(expr, permutation)
+        if isinstance(expr, ArrayTensorProduct):
+            expr, permutation = self._PermuteDims_denestarg_ArrayTensorProduct(expr, permutation)
+        if isinstance(expr, (ZeroArray, ZeroMatrix)):
+            return ZeroArray(*[expr.shape[i] for i in permutation.array_form])
+        plist = permutation.array_form
+        if plist == sorted(plist):
+            return expr
+        return self.func(expr, permutation, canonicalize=False)
+
+    @property
+    def expr(self):
+        return self.args[0]
+
+    @property
+    def permutation(self):
+        return self.args[1]
+
+    @classmethod
+    def _PermuteDims_denestarg_ArrayTensorProduct(cls, expr, permutation):
+        # Get the permutation in its image-form:
+        perm_image_form = _af_invert(permutation.array_form)
+        args = list(expr.args)
+        # Starting index global position for every arg:
+        cumul = list(accumulate([0] + expr.subranks))
+        # Split `perm_image_form` into a list of list corresponding to the indices
+        # of every argument:
+        perm_image_form_in_components = [perm_image_form[cumul[i]:cumul[i+1]] for i in range(len(args))]
+        # Create an index, target-position-key array:
+        ps = [(i, sorted(comp)) for i, comp in enumerate(perm_image_form_in_components)]
+        # Sort the array according to the target-position-key:
+        # In this way, we define a canonical way to sort the arguments according
+        # to the permutation.
+        ps.sort(key=lambda x: x[1])
+        # Read the inverse-permutation (i.e. image-form) of the args:
+        perm_args_image_form = [i[0] for i in ps]
+        # Apply the args-permutation to the `args`:
+        args_sorted = [args[i] for i in perm_args_image_form]
+        # Apply the args-permutation to the array-form of the permutation of the axes (of `expr`):
+        perm_image_form_sorted_args = [perm_image_form_in_components[i] for i in perm_args_image_form]
+        new_permutation = Permutation(_af_invert([j for i in perm_image_form_sorted_args for j in i]))
+        return _array_tensor_product(*args_sorted), new_permutation
+
+    @classmethod
+    def _PermuteDims_denestarg_ArrayContraction(cls, expr, permutation):
+        if not isinstance(expr, ArrayContraction):
+            return expr, permutation
+        if not isinstance(expr.expr, ArrayTensorProduct):
+            return expr, permutation
+        args = expr.expr.args
+        subranks = [get_rank(arg) for arg in expr.expr.args]
+
+        contraction_indices = expr.contraction_indices
+        contraction_indices_flat = [j for i in contraction_indices for j in i]
+        cumul = list(accumulate([0] + subranks))
+
+        # Spread the permutation in its array form across the args in the corresponding
+        # tensor-product arguments with free indices:
+        permutation_array_blocks_up = []
+        image_form = _af_invert(permutation.array_form)
+        counter = 0
+        for i, e in enumerate(subranks):
+            current = []
+            for j in range(cumul[i], cumul[i+1]):
+                if j in contraction_indices_flat:
+                    continue
+                current.append(image_form[counter])
+                counter += 1
+            permutation_array_blocks_up.append(current)
+
+        # Get the map of axis repositioning for every argument of tensor-product:
+        index_blocks = [list(range(cumul[i], cumul[i+1])) for i, e in enumerate(expr.subranks)]
+        index_blocks_up = expr._push_indices_up(expr.contraction_indices, index_blocks)
+        inverse_permutation = permutation**(-1)
+        index_blocks_up_permuted = [[inverse_permutation(j) for j in i if j is not None] for i in index_blocks_up]
+
+        # Sorting key is a list of tuple, first element is the index of `args`, second element of
+        # the tuple is the sorting key to sort `args` of the tensor product:
+        sorting_keys = list(enumerate(index_blocks_up_permuted))
+        sorting_keys.sort(key=lambda x: x[1])
+
+        # Now we can get the permutation acting on the args in its image-form:
+        new_perm_image_form = [i[0] for i in sorting_keys]
+        # Apply the args-level permutation to various elements:
+        new_index_blocks = [index_blocks[i] for i in new_perm_image_form]
+        new_index_perm_array_form = _af_invert([j for i in new_index_blocks for j in i])
+        new_args = [args[i] for i in new_perm_image_form]
+        new_contraction_indices = [tuple(new_index_perm_array_form[j] for j in i) for i in contraction_indices]
+        new_expr = _array_contraction(_array_tensor_product(*new_args), *new_contraction_indices)
+        new_permutation = Permutation(_af_invert([j for i in [permutation_array_blocks_up[k] for k in new_perm_image_form] for j in i]))
+        return new_expr, new_permutation
+
+    @classmethod
+    def _check_permutation_mapping(cls, expr, permutation):
+        subranks = expr.subranks
+        index2arg = [i for i, arg in enumerate(expr.args) for j in range(expr.subranks[i])]
+        permuted_indices = [permutation(i) for i in range(expr.subrank())]
+        new_args = list(expr.args)
+        arg_candidate_index = index2arg[permuted_indices[0]]
+        current_indices = []
+        new_permutation = []
+        inserted_arg_cand_indices = set()
+        for i, idx in enumerate(permuted_indices):
+            if index2arg[idx] != arg_candidate_index:
+                new_permutation.extend(current_indices)
+                current_indices = []
+                arg_candidate_index = index2arg[idx]
+            current_indices.append(idx)
+            arg_candidate_rank = subranks[arg_candidate_index]
+            if len(current_indices) == arg_candidate_rank:
+                new_permutation.extend(sorted(current_indices))
+                local_current_indices = [j - min(current_indices) for j in current_indices]
+                i1 = index2arg[i]
+                new_args[i1] = _permute_dims(new_args[i1], Permutation(local_current_indices))
+                inserted_arg_cand_indices.add(arg_candidate_index)
+                current_indices = []
+        new_permutation.extend(current_indices)
+
+        # TODO: swap args positions in order to simplify the expression:
+        # TODO: this should be in a function
+        args_positions = list(range(len(new_args)))
+        # Get possible shifts:
+        maps = {}
+        cumulative_subranks = [0] + list(accumulate(subranks))
+        for i in range(len(subranks)):
+            s = {index2arg[new_permutation[j]] for j in range(cumulative_subranks[i], cumulative_subranks[i+1])}
+            if len(s) != 1:
+                continue
+            elem = next(iter(s))
+            if i != elem:
+                maps[i] = elem
+
+        # Find cycles in the map:
+        lines = []
+        current_line = []
+        while maps:
+            if len(current_line) == 0:
+                k, v = maps.popitem()
+                current_line.append(k)
+            else:
+                k = current_line[-1]
+                if k not in maps:
+                    current_line = []
+                    continue
+                v = maps.pop(k)
+            if v in current_line:
+                lines.append(current_line)
+                current_line = []
+                continue
+            current_line.append(v)
+        for line in lines:
+            for i, e in enumerate(line):
+                args_positions[line[(i + 1) % len(line)]] = e
+
+        # TODO: function in order to permute the args:
+        permutation_blocks = [[new_permutation[cumulative_subranks[i] + j] for j in range(e)] for i, e in enumerate(subranks)]
+        new_args = [new_args[i] for i in args_positions]
+        new_permutation_blocks = [permutation_blocks[i] for i in args_positions]
+        new_permutation2 = [j for i in new_permutation_blocks for j in i]
+        return _array_tensor_product(*new_args), Permutation(new_permutation2)  # **(-1)
+
+    @classmethod
+    def _check_if_there_are_closed_cycles(cls, expr, permutation):
+        args = list(expr.args)
+        subranks = expr.subranks
+        cyclic_form = permutation.cyclic_form
+        cumulative_subranks = [0] + list(accumulate(subranks))
+        cyclic_min = [min(i) for i in cyclic_form]
+        cyclic_max = [max(i) for i in cyclic_form]
+        cyclic_keep = []
+        for i, cycle in enumerate(cyclic_form):
+            flag = True
+            for j in range(len(cumulative_subranks) - 1):
+                if cyclic_min[i] >= cumulative_subranks[j] and cyclic_max[i] < cumulative_subranks[j+1]:
+                    # Found a sinkable cycle.
+                    args[j] = _permute_dims(args[j], Permutation([[k - cumulative_subranks[j] for k in cyclic_form[i]]]))
+                    flag = False
+                    break
+            if flag:
+                cyclic_keep.append(cyclic_form[i])
+        return _array_tensor_product(*args), Permutation(cyclic_keep, size=permutation.size)
+
+    def nest_permutation(self):
+        r"""
+        DEPRECATED.
+        """
+        ret = self._nest_permutation(self.expr, self.permutation)
+        if ret is None:
+            return self
+        return ret
+
+    @classmethod
+    def _nest_permutation(cls, expr, permutation):
+        if isinstance(expr, ArrayTensorProduct):
+            return _permute_dims(*cls._check_if_there_are_closed_cycles(expr, permutation))
+        elif isinstance(expr, ArrayContraction):
+            # Invert tree hierarchy: put the contraction above.
+            cycles = permutation.cyclic_form
+            newcycles = ArrayContraction._convert_outer_indices_to_inner_indices(expr, *cycles)
+            newpermutation = Permutation(newcycles)
+            new_contr_indices = [tuple(newpermutation(j) for j in i) for i in expr.contraction_indices]
+            return _array_contraction(PermuteDims(expr.expr, newpermutation), *new_contr_indices)
+        elif isinstance(expr, ArrayAdd):
+            return _array_add(*[PermuteDims(arg, permutation) for arg in expr.args])
+        return None
+
+    def as_explicit(self):
+        expr = self.expr
+        if hasattr(expr, "as_explicit"):
+            expr = expr.as_explicit()
+        return permutedims(expr, self.permutation)
+
+    @classmethod
+    def _get_permutation_from_arguments(cls, permutation, index_order_old, index_order_new, dim):
+        if permutation is None:
+            if index_order_new is None or index_order_old is None:
+                raise ValueError("Permutation not defined")
+            return PermuteDims._get_permutation_from_index_orders(index_order_old, index_order_new, dim)
+        else:
+            if index_order_new is not None:
+                raise ValueError("index_order_new cannot be defined with permutation")
+            if index_order_old is not None:
+                raise ValueError("index_order_old cannot be defined with permutation")
+            return permutation
+
+    @classmethod
+    def _get_permutation_from_index_orders(cls, index_order_old, index_order_new, dim):
+        if len(set(index_order_new)) != dim:
+            raise ValueError("wrong number of indices in index_order_new")
+        if len(set(index_order_old)) != dim:
+            raise ValueError("wrong number of indices in index_order_old")
+        if len(set.symmetric_difference(set(index_order_new), set(index_order_old))) > 0:
+            raise ValueError("index_order_new and index_order_old must have the same indices")
+        permutation = [index_order_old.index(i) for i in index_order_new]
+        return permutation
+
+
+class ArrayDiagonal(_CodegenArrayAbstract):
+    r"""
+    Class to represent the diagonal operator.
+
+    Explanation
+    ===========
+
+    In a 2-dimensional array it returns the diagonal, this looks like the
+    operation:
+
+    `A_{ij} \rightarrow A_{ii}`
+
+    The diagonal over axes 1 and 2 (the second and third) of the tensor product
+    of two 2-dimensional arrays `A \otimes B` is
+
+    `\Big[ A_{ab} B_{cd} \Big]_{abcd} \rightarrow \Big[ A_{ai} B_{id} \Big]_{adi}`
+
+    In this last example the array expression has been reduced from
+    4-dimensional to 3-dimensional. Notice that no contraction has occurred,
+    rather there is a new index `i` for the diagonal, contraction would have
+    reduced the array to 2 dimensions.
+
+    Notice that the diagonalized out dimensions are added as new dimensions at
+    the end of the indices.
+    """
+
+    def __new__(cls, expr, *diagonal_indices, **kwargs):
+        expr = _sympify(expr)
+        diagonal_indices = [Tuple(*sorted(i)) for i in diagonal_indices]
+        canonicalize = kwargs.get("canonicalize", False)
+
+        shape = get_shape(expr)
+        if shape is not None:
+            cls._validate(expr, *diagonal_indices, **kwargs)
+            # Get new shape:
+            positions, shape = cls._get_positions_shape(shape, diagonal_indices)
+        else:
+            positions = None
+        if len(diagonal_indices) == 0:
+            return expr
+        obj = Basic.__new__(cls, expr, *diagonal_indices)
+        obj._positions = positions
+        obj._subranks = _get_subranks(expr)
+        obj._shape = shape
+        if canonicalize:
+            return obj._canonicalize()
+        return obj
+
+    def _canonicalize(self):
+        expr = self.expr
+        diagonal_indices = self.diagonal_indices
+        trivial_diags = [i for i in diagonal_indices if len(i) == 1]
+        if len(trivial_diags) > 0:
+            trivial_pos = {e[0]: i for i, e in enumerate(diagonal_indices) if len(e) == 1}
+            diag_pos = {e: i for i, e in enumerate(diagonal_indices) if len(e) > 1}
+            diagonal_indices_short = [i for i in diagonal_indices if len(i) > 1]
+            rank1 = get_rank(self)
+            rank2 = len(diagonal_indices)
+            rank3 = rank1 - rank2
+            inv_permutation = []
+            counter1 = 0
+            indices_down = ArrayDiagonal._push_indices_down(diagonal_indices_short, list(range(rank1)), get_rank(expr))
+            for i in indices_down:
+                if i in trivial_pos:
+                    inv_permutation.append(rank3 + trivial_pos[i])
+                elif isinstance(i, (Integer, int)):
+                    inv_permutation.append(counter1)
+                    counter1 += 1
+                else:
+                    inv_permutation.append(rank3 + diag_pos[i])
+            permutation = _af_invert(inv_permutation)
+            if len(diagonal_indices_short) > 0:
+                return _permute_dims(_array_diagonal(expr, *diagonal_indices_short), permutation)
+            else:
+                return _permute_dims(expr, permutation)
+        if isinstance(expr, ArrayAdd):
+            return self._ArrayDiagonal_denest_ArrayAdd(expr, *diagonal_indices)
+        if isinstance(expr, ArrayDiagonal):
+            return self._ArrayDiagonal_denest_ArrayDiagonal(expr, *diagonal_indices)
+        if isinstance(expr, PermuteDims):
+            return self._ArrayDiagonal_denest_PermuteDims(expr, *diagonal_indices)
+        if isinstance(expr, (ZeroArray, ZeroMatrix)):
+            positions, shape = self._get_positions_shape(expr.shape, diagonal_indices)
+            return ZeroArray(*shape)
+        return self.func(expr, *diagonal_indices, canonicalize=False)
+
+    @staticmethod
+    def _validate(expr, *diagonal_indices, **kwargs):
+        # Check that no diagonalization happens on indices with mismatched
+        # dimensions:
+        shape = get_shape(expr)
+        for i in diagonal_indices:
+            if any(j >= len(shape) for j in i):
+                raise ValueError("index is larger than expression shape")
+            if len({shape[j] for j in i}) != 1:
+                raise ValueError("diagonalizing indices of different dimensions")
+            if not kwargs.get("allow_trivial_diags", False) and len(i) <= 1:
+                raise ValueError("need at least two axes to diagonalize")
+            if len(set(i)) != len(i):
+                raise ValueError("axis index cannot be repeated")
+
+    @staticmethod
+    def _remove_trivial_dimensions(shape, *diagonal_indices):
+        return [tuple(j for j in i) for i in diagonal_indices if shape[i[0]] != 1]
+
+    @property
+    def expr(self):
+        return self.args[0]
+
+    @property
+    def diagonal_indices(self):
+        return self.args[1:]
+
+    @staticmethod
+    def _flatten(expr, *outer_diagonal_indices):
+        inner_diagonal_indices = expr.diagonal_indices
+        all_inner = [j for i in inner_diagonal_indices for j in i]
+        all_inner.sort()
+        # TODO: add API for total rank and cumulative rank:
+        total_rank = _get_subrank(expr)
+        inner_rank = len(all_inner)
+        outer_rank = total_rank - inner_rank
+        shifts = [0 for i in range(outer_rank)]
+        counter = 0
+        pointer = 0
+        for i in range(outer_rank):
+            while pointer < inner_rank and counter >= all_inner[pointer]:
+                counter += 1
+                pointer += 1
+            shifts[i] += pointer
+            counter += 1
+        outer_diagonal_indices = tuple(tuple(shifts[j] + j for j in i) for i in outer_diagonal_indices)
+        diagonal_indices = inner_diagonal_indices + outer_diagonal_indices
+        return _array_diagonal(expr.expr, *diagonal_indices)
+
+    @classmethod
+    def _ArrayDiagonal_denest_ArrayAdd(cls, expr, *diagonal_indices):
+        return _array_add(*[_array_diagonal(arg, *diagonal_indices) for arg in expr.args])
+
+    @classmethod
+    def _ArrayDiagonal_denest_ArrayDiagonal(cls, expr, *diagonal_indices):
+        return cls._flatten(expr, *diagonal_indices)
+
+    @classmethod
+    def _ArrayDiagonal_denest_PermuteDims(cls, expr: PermuteDims, *diagonal_indices):
+        back_diagonal_indices = [[expr.permutation(j) for j in i] for i in diagonal_indices]
+        nondiag = [i for i in range(get_rank(expr)) if not any(i in j for j in diagonal_indices)]
+        back_nondiag = [expr.permutation(i) for i in nondiag]
+        remap = {e: i for i, e in enumerate(sorted(back_nondiag))}
+        new_permutation1 = [remap[i] for i in back_nondiag]
+        shift = len(new_permutation1)
+        diag_block_perm = [i + shift for i in range(len(back_diagonal_indices))]
+        new_permutation = new_permutation1 + diag_block_perm
+        return _permute_dims(
+            _array_diagonal(
+                expr.expr,
+                *back_diagonal_indices
+            ),
+            new_permutation
+        )
+
+    def _push_indices_down_nonstatic(self, indices):
+        transform = lambda x: self._positions[x] if x < len(self._positions) else None
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    def _push_indices_up_nonstatic(self, indices):
+
+        def transform(x):
+            for i, e in enumerate(self._positions):
+                if (isinstance(e, int) and x == e) or (isinstance(e, tuple) and x in e):
+                    return i
+
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    @classmethod
+    def _push_indices_down(cls, diagonal_indices, indices, rank):
+        positions, shape = cls._get_positions_shape(range(rank), diagonal_indices)
+        transform = lambda x: positions[x] if x < len(positions) else None
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    @classmethod
+    def _push_indices_up(cls, diagonal_indices, indices, rank):
+        positions, shape = cls._get_positions_shape(range(rank), diagonal_indices)
+
+        def transform(x):
+            for i, e in enumerate(positions):
+                if (isinstance(e, int) and x == e) or (isinstance(e, (tuple, Tuple)) and (x in e)):
+                    return i
+
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    @classmethod
+    def _get_positions_shape(cls, shape, diagonal_indices):
+        data1 = tuple((i, shp) for i, shp in enumerate(shape) if not any(i in j for j in diagonal_indices))
+        pos1, shp1 = zip(*data1) if data1 else ((), ())
+        data2 = tuple((i, shape[i[0]]) for i in diagonal_indices)
+        pos2, shp2 = zip(*data2) if data2 else ((), ())
+        positions = pos1 + pos2
+        shape = shp1 + shp2
+        return positions, shape
+
+    def as_explicit(self):
+        expr = self.expr
+        if hasattr(expr, "as_explicit"):
+            expr = expr.as_explicit()
+        return tensordiagonal(expr, *self.diagonal_indices)
+
+
+class ArrayElementwiseApplyFunc(_CodegenArrayAbstract):
+
+    def __new__(cls, function, element):
+
+        if not isinstance(function, Lambda):
+            d = Dummy('d')
+            function = Lambda(d, function(d))
+
+        obj = _CodegenArrayAbstract.__new__(cls, function, element)
+        obj._subranks = _get_subranks(element)
+        return obj
+
+    @property
+    def function(self):
+        return self.args[0]
+
+    @property
+    def expr(self):
+        return self.args[1]
+
+    @property
+    def shape(self):
+        return self.expr.shape
+
+    def _get_function_fdiff(self):
+        d = Dummy("d")
+        function = self.function(d)
+        fdiff = function.diff(d)
+        if isinstance(fdiff, Function):
+            fdiff = type(fdiff)
+        else:
+            fdiff = Lambda(d, fdiff)
+        return fdiff
+
+    def as_explicit(self):
+        expr = self.expr
+        if hasattr(expr, "as_explicit"):
+            expr = expr.as_explicit()
+        return expr.applyfunc(self.function)
+
+
+class ArrayContraction(_CodegenArrayAbstract):
+    r"""
+    This class is meant to represent contractions of arrays in a form easily
+    processable by the code printers.
+    """
+
+    def __new__(cls, expr, *contraction_indices, **kwargs):
+        contraction_indices = _sort_contraction_indices(contraction_indices)
+        expr = _sympify(expr)
+
+        canonicalize = kwargs.get("canonicalize", False)
+
+        obj = Basic.__new__(cls, expr, *contraction_indices)
+        obj._subranks = _get_subranks(expr)
+        obj._mapping = _get_mapping_from_subranks(obj._subranks)
+
+        free_indices_to_position = {i: i for i in range(sum(obj._subranks)) if all(i not in cind for cind in contraction_indices)}
+        obj._free_indices_to_position = free_indices_to_position
+
+        shape = get_shape(expr)
+        cls._validate(expr, *contraction_indices)
+        if shape:
+            shape = tuple(shp for i, shp in enumerate(shape) if not any(i in j for j in contraction_indices))
+        obj._shape = shape
+        if canonicalize:
+            return obj._canonicalize()
+        return obj
+
+    def _canonicalize(self):
+        expr = self.expr
+        contraction_indices = self.contraction_indices
+
+        if len(contraction_indices) == 0:
+            return expr
+
+        if isinstance(expr, ArrayContraction):
+            return self._ArrayContraction_denest_ArrayContraction(expr, *contraction_indices)
+
+        if isinstance(expr, (ZeroArray, ZeroMatrix)):
+            return self._ArrayContraction_denest_ZeroArray(expr, *contraction_indices)
+
+        if isinstance(expr, PermuteDims):
+            return self._ArrayContraction_denest_PermuteDims(expr, *contraction_indices)
+
+        if isinstance(expr, ArrayTensorProduct):
+            expr, contraction_indices = self._sort_fully_contracted_args(expr, contraction_indices)
+            expr, contraction_indices = self._lower_contraction_to_addends(expr, contraction_indices)
+            if len(contraction_indices) == 0:
+                return expr
+
+        if isinstance(expr, ArrayDiagonal):
+            return self._ArrayContraction_denest_ArrayDiagonal(expr, *contraction_indices)
+
+        if isinstance(expr, ArrayAdd):
+            return self._ArrayContraction_denest_ArrayAdd(expr, *contraction_indices)
+
+        # Check single index contractions on 1-dimensional axes:
+        contraction_indices = [i for i in contraction_indices if len(i) > 1 or get_shape(expr)[i[0]] != 1]
+        if len(contraction_indices) == 0:
+            return expr
+
+        return self.func(expr, *contraction_indices, canonicalize=False)
+
+    def __mul__(self, other):
+        if other == 1:
+            return self
+        else:
+            raise NotImplementedError("Product of N-dim arrays is not uniquely defined. Use another method.")
+
+    def __rmul__(self, other):
+        if other == 1:
+            return self
+        else:
+            raise NotImplementedError("Product of N-dim arrays is not uniquely defined. Use another method.")
+
+    @staticmethod
+    def _validate(expr, *contraction_indices):
+        shape = get_shape(expr)
+        if shape is None:
+            return
+
+        # Check that no contraction happens when the shape is mismatched:
+        for i in contraction_indices:
+            if len({shape[j] for j in i if shape[j] != -1}) != 1:
+                raise ValueError("contracting indices of different dimensions")
+
+    @classmethod
+    def _push_indices_down(cls, contraction_indices, indices):
+        flattened_contraction_indices = [j for i in contraction_indices for j in i]
+        flattened_contraction_indices.sort()
+        transform = _build_push_indices_down_func_transformation(flattened_contraction_indices)
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    @classmethod
+    def _push_indices_up(cls, contraction_indices, indices):
+        flattened_contraction_indices = [j for i in contraction_indices for j in i]
+        flattened_contraction_indices.sort()
+        transform = _build_push_indices_up_func_transformation(flattened_contraction_indices)
+        return _apply_recursively_over_nested_lists(transform, indices)
+
+    @classmethod
+    def _lower_contraction_to_addends(cls, expr, contraction_indices):
+        if isinstance(expr, ArrayAdd):
+            raise NotImplementedError()
+        if not isinstance(expr, ArrayTensorProduct):
+            return expr, contraction_indices
+        subranks = expr.subranks
+        cumranks = list(accumulate([0] + subranks))
+        contraction_indices_remaining = []
+        contraction_indices_args = [[] for i in expr.args]
+        backshift = set()
+        for contraction_group in contraction_indices:
+            for j in range(len(expr.args)):
+                if not isinstance(expr.args[j], ArrayAdd):
+                    continue
+                if all(cumranks[j] <= k < cumranks[j+1] for k in contraction_group):
+                    contraction_indices_args[j].append([k - cumranks[j] for k in contraction_group])
+                    backshift.update(contraction_group)
+                    break
+            else:
+                contraction_indices_remaining.append(contraction_group)
+        if len(contraction_indices_remaining) == len(contraction_indices):
+            return expr, contraction_indices
+        total_rank = get_rank(expr)
+        shifts = list(accumulate([1 if i in backshift else 0 for i in range(total_rank)]))
+        contraction_indices_remaining = [Tuple.fromiter(j - shifts[j] for j in i) for i in contraction_indices_remaining]
+        ret = _array_tensor_product(*[
+            _array_contraction(arg, *contr) for arg, contr in zip(expr.args, contraction_indices_args)
+        ])
+        return ret, contraction_indices_remaining
+
+    def split_multiple_contractions(self):
+        """
+        Recognize multiple contractions and attempt at rewriting them as paired-contractions.
+
+        This allows some contractions involving more than two indices to be
+        rewritten as multiple contractions involving two indices, thus allowing
+        the expression to be rewritten as a matrix multiplication line.
+
+        Examples:
+
+        * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C`
+
+        Care for:
+        - matrix being diagonalized (i.e. `A_ii`)
+        - vectors being diagonalized (i.e. `a_i0`)
+
+        Multiple contractions can be split into matrix multiplications if
+        not more than two arguments are non-diagonals or non-vectors.
+        Vectors get diagonalized while diagonal matrices remain diagonal.
+        The non-diagonal matrices can be at the beginning or at the end
+        of the final matrix multiplication line.
+        """
+
+        editor = _EditArrayContraction(self)
+
+        contraction_indices = self.contraction_indices
+
+        onearray_insert = []
+
+        for indl, links in enumerate(contraction_indices):
+            if len(links) <= 2:
+                continue
+
+            # Check multiple contractions:
+            #
+            # Examples:
+            #
+            # * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C \otimes OneArray(1)` with permutation (1 2)
+            #
+            # Care for:
+            # - matrix being diagonalized (i.e. `A_ii`)
+            # - vectors being diagonalized (i.e. `a_i0`)
+
+            # Multiple contractions can be split into matrix multiplications if
+            # not more than three arguments are non-diagonals or non-vectors.
+            #
+            # Vectors get diagonalized while diagonal matrices remain diagonal.
+            # The non-diagonal matrices can be at the beginning or at the end
+            # of the final matrix multiplication line.
+
+            positions = editor.get_mapping_for_index(indl)
+
+            # Also consider the case of diagonal matrices being contracted:
+            current_dimension = self.expr.shape[links[0]]
+
+            not_vectors = []
+            vectors = []
+            for arg_ind, rel_ind in positions:
+                arg = editor.args_with_ind[arg_ind]
+                mat = arg.element
+                abs_arg_start, abs_arg_end = editor.get_absolute_range(arg)
+                other_arg_pos = 1-rel_ind
+                other_arg_abs = abs_arg_start + other_arg_pos
+                if ((1 not in mat.shape) or
+                    ((current_dimension == 1) is True and mat.shape != (1, 1)) or
+                    any(other_arg_abs in l for li, l in enumerate(contraction_indices) if li != indl)
+                ):
+                    not_vectors.append((arg, rel_ind))
+                else:
+                    vectors.append((arg, rel_ind))
+            if len(not_vectors) > 2:
+                # If more than two arguments in the multiple contraction are
+                # non-vectors and non-diagonal matrices, we cannot find a way
+                # to split this contraction into a matrix multiplication line:
+                continue
+            # Three cases to handle:
+            # - zero non-vectors
+            # - one non-vector
+            # - two non-vectors
+            for v, rel_ind in vectors:
+                v.element = diagonalize_vector(v.element)
+            vectors_to_loop = not_vectors[:1] + vectors + not_vectors[1:]
+            first_not_vector, rel_ind = vectors_to_loop[0]
+            new_index = first_not_vector.indices[rel_ind]
+
+            for v, rel_ind in vectors_to_loop[1:-1]:
+                v.indices[rel_ind] = new_index
+                new_index = editor.get_new_contraction_index()
+                assert v.indices.index(None) == 1 - rel_ind
+                v.indices[v.indices.index(None)] = new_index
+                onearray_insert.append(v)
+
+            last_vec, rel_ind = vectors_to_loop[-1]
+            last_vec.indices[rel_ind] = new_index
+
+        for v in onearray_insert:
+            editor.insert_after(v, _ArgE(OneArray(1), [None]))
+
+        return editor.to_array_contraction()
+
+    def flatten_contraction_of_diagonal(self):
+        if not isinstance(self.expr, ArrayDiagonal):
+            return self
+        contraction_down = self.expr._push_indices_down(self.expr.diagonal_indices, self.contraction_indices)
+        new_contraction_indices = []
+        diagonal_indices = self.expr.diagonal_indices[:]
+        for i in contraction_down:
+            contraction_group = list(i)
+            for j in i:
+                diagonal_with = [k for k in diagonal_indices if j in k]
+                contraction_group.extend([l for k in diagonal_with for l in k])
+                diagonal_indices = [k for k in diagonal_indices if k not in diagonal_with]
+            new_contraction_indices.append(sorted(set(contraction_group)))
+
+        new_contraction_indices = ArrayDiagonal._push_indices_up(diagonal_indices, new_contraction_indices)
+        return _array_contraction(
+            _array_diagonal(
+                self.expr.expr,
+                *diagonal_indices
+            ),
+            *new_contraction_indices
+        )
+
+    @staticmethod
+    def _get_free_indices_to_position_map(free_indices, contraction_indices):
+        free_indices_to_position = {}
+        flattened_contraction_indices = [j for i in contraction_indices for j in i]
+        counter = 0
+        for ind in free_indices:
+            while counter in flattened_contraction_indices:
+                counter += 1
+            free_indices_to_position[ind] = counter
+            counter += 1
+        return free_indices_to_position
+
+    @staticmethod
+    def _get_index_shifts(expr):
+        """
+        Get the mapping of indices at the positions before the contraction
+        occurs.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array import tensorproduct, tensorcontraction
+        >>> from sympy import MatrixSymbol
+        >>> M = MatrixSymbol("M", 3, 3)
+        >>> N = MatrixSymbol("N", 3, 3)
+        >>> cg = tensorcontraction(tensorproduct(M, N), [1, 2])
+        >>> cg._get_index_shifts(cg)
+        [0, 2]
+
+        Indeed, ``cg`` after the contraction has two dimensions, 0 and 1. They
+        need to be shifted by 0 and 2 to get the corresponding positions before
+        the contraction (that is, 0 and 3).
+        """
+        inner_contraction_indices = expr.contraction_indices
+        all_inner = [j for i in inner_contraction_indices for j in i]
+        all_inner.sort()
+        # TODO: add API for total rank and cumulative rank:
+        total_rank = _get_subrank(expr)
+        inner_rank = len(all_inner)
+        outer_rank = total_rank - inner_rank
+        shifts = [0 for i in range(outer_rank)]
+        counter = 0
+        pointer = 0
+        for i in range(outer_rank):
+            while pointer < inner_rank and counter >= all_inner[pointer]:
+                counter += 1
+                pointer += 1
+            shifts[i] += pointer
+            counter += 1
+        return shifts
+
+    @staticmethod
+    def _convert_outer_indices_to_inner_indices(expr, *outer_contraction_indices):
+        shifts = ArrayContraction._get_index_shifts(expr)
+        outer_contraction_indices = tuple(tuple(shifts[j] + j for j in i) for i in outer_contraction_indices)
+        return outer_contraction_indices
+
+    @staticmethod
+    def _flatten(expr, *outer_contraction_indices):
+        inner_contraction_indices = expr.contraction_indices
+        outer_contraction_indices = ArrayContraction._convert_outer_indices_to_inner_indices(expr, *outer_contraction_indices)
+        contraction_indices = inner_contraction_indices + outer_contraction_indices
+        return _array_contraction(expr.expr, *contraction_indices)
+
+    @classmethod
+    def _ArrayContraction_denest_ArrayContraction(cls, expr, *contraction_indices):
+        return cls._flatten(expr, *contraction_indices)
+
+    @classmethod
+    def _ArrayContraction_denest_ZeroArray(cls, expr, *contraction_indices):
+        contraction_indices_flat = [j for i in contraction_indices for j in i]
+        shape = [e for i, e in enumerate(expr.shape) if i not in contraction_indices_flat]
+        return ZeroArray(*shape)
+
+    @classmethod
+    def _ArrayContraction_denest_ArrayAdd(cls, expr, *contraction_indices):
+        return _array_add(*[_array_contraction(i, *contraction_indices) for i in expr.args])
+
+    @classmethod
+    def _ArrayContraction_denest_PermuteDims(cls, expr, *contraction_indices):
+        permutation = expr.permutation
+        plist = permutation.array_form
+        new_contraction_indices = [tuple(permutation(j) for j in i) for i in contraction_indices]
+        new_plist = [i for i in plist if not any(i in j for j in new_contraction_indices)]
+        new_plist = cls._push_indices_up(new_contraction_indices, new_plist)
+        return _permute_dims(
+            _array_contraction(expr.expr, *new_contraction_indices),
+            Permutation(new_plist)
+        )
+
+    @classmethod
+    def _ArrayContraction_denest_ArrayDiagonal(cls, expr: 'ArrayDiagonal', *contraction_indices):
+        diagonal_indices = list(expr.diagonal_indices)
+        down_contraction_indices = expr._push_indices_down(expr.diagonal_indices, contraction_indices, get_rank(expr.expr))
+        # Flatten diagonally contracted indices:
+        down_contraction_indices = [[k for j in i for k in (j if isinstance(j, (tuple, Tuple)) else [j])] for i in down_contraction_indices]
+        new_contraction_indices = []
+        for contr_indgrp in down_contraction_indices:
+            ind = contr_indgrp[:]
+            for j, diag_indgrp in enumerate(diagonal_indices):
+                if diag_indgrp is None:
+                    continue
+                if any(i in diag_indgrp for i in contr_indgrp):
+                    ind.extend(diag_indgrp)
+                    diagonal_indices[j] = None
+            new_contraction_indices.append(sorted(set(ind)))
+
+        new_diagonal_indices_down = [i for i in diagonal_indices if i is not None]
+        new_diagonal_indices = ArrayContraction._push_indices_up(new_contraction_indices, new_diagonal_indices_down)
+        return _array_diagonal(
+            _array_contraction(expr.expr, *new_contraction_indices),
+            *new_diagonal_indices
+        )
+
+    @classmethod
+    def _sort_fully_contracted_args(cls, expr, contraction_indices):
+        if expr.shape is None:
+            return expr, contraction_indices
+        cumul = list(accumulate([0] + expr.subranks))
+        index_blocks = [list(range(cumul[i], cumul[i+1])) for i in range(len(expr.args))]
+        contraction_indices_flat = {j for i in contraction_indices for j in i}
+        fully_contracted = [all(j in contraction_indices_flat for j in range(cumul[i], cumul[i+1])) for i, arg in enumerate(expr.args)]
+        new_pos = sorted(range(len(expr.args)), key=lambda x: (0, default_sort_key(expr.args[x])) if fully_contracted[x] else (1,))
+        new_args = [expr.args[i] for i in new_pos]
+        new_index_blocks_flat = [j for i in new_pos for j in index_blocks[i]]
+        index_permutation_array_form = _af_invert(new_index_blocks_flat)
+        new_contraction_indices = [tuple(index_permutation_array_form[j] for j in i) for i in contraction_indices]
+        new_contraction_indices = _sort_contraction_indices(new_contraction_indices)
+        return _array_tensor_product(*new_args), new_contraction_indices
+
+    def _get_contraction_tuples(self):
+        r"""
+        Return tuples containing the argument index and position within the
+        argument of the index position.
+
+        Examples
+        ========
+
+        >>> from sympy import MatrixSymbol
+        >>> from sympy.abc import N
+        >>> from sympy.tensor.array import tensorproduct, tensorcontraction
+        >>> A = MatrixSymbol("A", N, N)
+        >>> B = MatrixSymbol("B", N, N)
+
+        >>> cg = tensorcontraction(tensorproduct(A, B), (1, 2))
+        >>> cg._get_contraction_tuples()
+        [[(0, 1), (1, 0)]]
+
+        Notes
+        =====
+
+        Here the contraction pair `(1, 2)` meaning that the 2nd and 3rd indices
+        of the tensor product `A\otimes B` are contracted, has been transformed
+        into `(0, 1)` and `(1, 0)`, identifying the same indices in a different
+        notation. `(0, 1)` is the second index (1) of the first argument (i.e.
+                0 or `A`). `(1, 0)` is the first index (i.e. 0) of the second
+        argument (i.e. 1 or `B`).
+        """
+        mapping = self._mapping
+        return [[mapping[j] for j in i] for i in self.contraction_indices]
+
+    @staticmethod
+    def _contraction_tuples_to_contraction_indices(expr, contraction_tuples):
+        # TODO: check that `expr` has `.subranks`:
+        ranks = expr.subranks
+        cumulative_ranks = [0] + list(accumulate(ranks))
+        return [tuple(cumulative_ranks[j]+k for j, k in i) for i in contraction_tuples]
+
+    @property
+    def free_indices(self):
+        return self._free_indices[:]
+
+    @property
+    def free_indices_to_position(self):
+        return dict(self._free_indices_to_position)
+
+    @property
+    def expr(self):
+        return self.args[0]
+
+    @property
+    def contraction_indices(self):
+        return self.args[1:]
+
+    def _contraction_indices_to_components(self):
+        expr = self.expr
+        if not isinstance(expr, ArrayTensorProduct):
+            raise NotImplementedError("only for contractions of tensor products")
+        ranks = expr.subranks
+        mapping = {}
+        counter = 0
+        for i, rank in enumerate(ranks):
+            for j in range(rank):
+                mapping[counter] = (i, j)
+                counter += 1
+        return mapping
+
+    def sort_args_by_name(self):
+        """
+        Sort arguments in the tensor product so that their order is lexicographical.
+
+        Examples
+        ========
+
+        >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+        >>> from sympy import MatrixSymbol
+        >>> from sympy.abc import N
+        >>> A = MatrixSymbol("A", N, N)
+        >>> B = MatrixSymbol("B", N, N)
+        >>> C = MatrixSymbol("C", N, N)
+        >>> D = MatrixSymbol("D", N, N)
+
+        >>> cg = convert_matrix_to_array(C*D*A*B)
+        >>> cg
+        ArrayContraction(ArrayTensorProduct(A, D, C, B), (0, 3), (1, 6), (2, 5))
+        >>> cg.sort_args_by_name()
+        ArrayContraction(ArrayTensorProduct(A, D, B, C), (0, 3), (1, 4), (2, 7))
+        """
+        expr = self.expr
+        if not isinstance(expr, ArrayTensorProduct):
+            return self
+        args = expr.args
+        sorted_data = sorted(enumerate(args), key=lambda x: default_sort_key(x[1]))
+        pos_sorted, args_sorted = zip(*sorted_data)
+        reordering_map = {i: pos_sorted.index(i) for i, arg in enumerate(args)}
+        contraction_tuples = self._get_contraction_tuples()
+        contraction_tuples = [[(reordering_map[j], k) for j, k in i] for i in contraction_tuples]
+        c_tp = _array_tensor_product(*args_sorted)
+        new_contr_indices = self._contraction_tuples_to_contraction_indices(
+                c_tp,
+                contraction_tuples
+        )
+        return _array_contraction(c_tp, *new_contr_indices)
+
+    def _get_contraction_links(self):
+        r"""
+        Returns a dictionary of links between arguments in the tensor product
+        being contracted.
+
+        See the example for an explanation of the values.
+
+        Examples
+        ========
+
+        >>> from sympy import MatrixSymbol
+        >>> from sympy.abc import N
+        >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+        >>> A = MatrixSymbol("A", N, N)
+        >>> B = MatrixSymbol("B", N, N)
+        >>> C = MatrixSymbol("C", N, N)
+        >>> D = MatrixSymbol("D", N, N)
+
+        Matrix multiplications are pairwise contractions between neighboring
+        matrices:
+
+        `A_{ij} B_{jk} C_{kl} D_{lm}`
+
+        >>> cg = convert_matrix_to_array(A*B*C*D)
+        >>> cg
+        ArrayContraction(ArrayTensorProduct(B, C, A, D), (0, 5), (1, 2), (3, 6))
+
+        >>> cg._get_contraction_links()
+        {0: {0: (2, 1), 1: (1, 0)}, 1: {0: (0, 1), 1: (3, 0)}, 2: {1: (0, 0)}, 3: {0: (1, 1)}}
+
+        This dictionary is interpreted as follows: argument in position 0 (i.e.
+        matrix `A`) has its second index (i.e. 1) contracted to `(1, 0)`, that
+        is argument in position 1 (matrix `B`) on the first index slot of `B`,
+        this is the contraction provided by the index `j` from `A`.
+
+        The argument in position 1 (that is, matrix `B`) has two contractions,
+        the ones provided by the indices `j` and `k`, respectively the first
+        and second indices (0 and 1 in the sub-dict).  The link `(0, 1)` and
+        `(2, 0)` respectively. `(0, 1)` is the index slot 1 (the 2nd) of
+        argument in position 0 (that is, `A_{\ldot j}`), and so on.
+        """
+        args, dlinks = _get_contraction_links([self], self.subranks, *self.contraction_indices)
+        return dlinks
+
+    def as_explicit(self):
+        expr = self.expr
+        if hasattr(expr, "as_explicit"):
+            expr = expr.as_explicit()
+        return tensorcontraction(expr, *self.contraction_indices)
+
+
+class Reshape(_CodegenArrayAbstract):
+    """
+    Reshape the dimensions of an array expression.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array.expressions import ArraySymbol, Reshape
+    >>> A = ArraySymbol("A", (6,))
+    >>> A.shape
+    (6,)
+    >>> Reshape(A, (3, 2)).shape
+    (3, 2)
+
+    Check the component-explicit forms:
+
+    >>> A.as_explicit()
+    [A[0], A[1], A[2], A[3], A[4], A[5]]
+    >>> Reshape(A, (3, 2)).as_explicit()
+    [[A[0], A[1]], [A[2], A[3]], [A[4], A[5]]]
+
+    """
+
+    def __new__(cls, expr, shape):
+        expr = _sympify(expr)
+        if not isinstance(shape, Tuple):
+            shape = Tuple(*shape)
+        if Equality(Mul.fromiter(expr.shape), Mul.fromiter(shape)) == False:
+            raise ValueError("shape mismatch")
+        obj = Expr.__new__(cls, expr, shape)
+        obj._shape = tuple(shape)
+        obj._expr = expr
+        return obj
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def expr(self):
+        return self._expr
+
+    def doit(self, *args, **kwargs):
+        if kwargs.get("deep", True):
+            expr = self.expr.doit(*args, **kwargs)
+        else:
+            expr = self.expr
+        if isinstance(expr, (MatrixBase, NDimArray)):
+            return expr.reshape(*self.shape)
+        return Reshape(expr, self.shape)
+
+    def as_explicit(self):
+        ee = self.expr
+        if hasattr(ee, "as_explicit"):
+            ee = ee.as_explicit()
+        if isinstance(ee, MatrixBase):
+            from sympy import Array
+            ee = Array(ee)
+        elif isinstance(ee, MatrixExpr):
+            return self
+        return ee.reshape(*self.shape)
+
+
+class _ArgE:
+    """
+    The ``_ArgE`` object contains references to the array expression
+    (``.element``) and a list containing the information about index
+    contractions (``.indices``).
+
+    Index contractions are numbered and contracted indices show the number of
+    the contraction. Uncontracted indices have ``None`` value.
+
+    For example:
+    ``_ArgE(M, [None, 3])``
+    This object means that expression ``M`` is part of an array contraction
+    and has two indices, the first is not contracted (value ``None``),
+    the second index is contracted to the 4th (i.e. number ``3``) group of the
+    array contraction object.
+    """
+    indices: List[Optional[int]]
+
+    def __init__(self, element, indices: Optional[List[Optional[int]]] = None):
+        self.element = element
+        if indices is None:
+            self.indices = [None for i in range(get_rank(element))]
+        else:
+            self.indices = indices
+
+    def __str__(self):
+        return "_ArgE(%s, %s)" % (self.element, self.indices)
+
+    __repr__ = __str__
+
+
+class _IndPos:
+    """
+    Index position, requiring two integers in the constructor:
+
+    - arg: the position of the argument in the tensor product,
+    - rel: the relative position of the index inside the argument.
+    """
+    def __init__(self, arg: int, rel: int):
+        self.arg = arg
+        self.rel = rel
+
+    def __str__(self):
+        return "_IndPos(%i, %i)" % (self.arg, self.rel)
+
+    __repr__ = __str__
+
+    def __iter__(self):
+        yield from [self.arg, self.rel]
+
+
+class _EditArrayContraction:
+    """
+    Utility class to help manipulate array contraction objects.
+
+    This class takes as input an ``ArrayContraction`` object and turns it into
+    an editable object.
+
+    The field ``args_with_ind`` of this class is a list of ``_ArgE`` objects
+    which can be used to easily edit the contraction structure of the
+    expression.
+
+    Once editing is finished, the ``ArrayContraction`` object may be recreated
+    by calling the ``.to_array_contraction()`` method.
+    """
+
+    def __init__(self, base_array: typing.Union[ArrayContraction, ArrayDiagonal, ArrayTensorProduct]):
+
+        expr: Basic
+        diagonalized: tTuple[tTuple[int, ...], ...]
+        contraction_indices: List[tTuple[int]]
+        if isinstance(base_array, ArrayContraction):
+            mapping = _get_mapping_from_subranks(base_array.subranks)
+            expr = base_array.expr
+            contraction_indices = base_array.contraction_indices
+            diagonalized = ()
+        elif isinstance(base_array, ArrayDiagonal):
+
+            if isinstance(base_array.expr, ArrayContraction):
+                mapping = _get_mapping_from_subranks(base_array.expr.subranks)
+                expr = base_array.expr.expr
+                diagonalized = ArrayContraction._push_indices_down(base_array.expr.contraction_indices, base_array.diagonal_indices)
+                contraction_indices = base_array.expr.contraction_indices
+            elif isinstance(base_array.expr, ArrayTensorProduct):
+                mapping = {}
+                expr = base_array.expr
+                diagonalized = base_array.diagonal_indices
+                contraction_indices = []
+            else:
+                mapping = {}
+                expr = base_array.expr
+                diagonalized = base_array.diagonal_indices
+                contraction_indices = []
+
+        elif isinstance(base_array, ArrayTensorProduct):
+            expr = base_array
+            contraction_indices = []
+            diagonalized = ()
+        else:
+            raise NotImplementedError()
+
+        if isinstance(expr, ArrayTensorProduct):
+            args = list(expr.args)
+        else:
+            args = [expr]
+
+        args_with_ind: List[_ArgE] = [_ArgE(arg) for arg in args]
+        for i, contraction_tuple in enumerate(contraction_indices):
+            for j in contraction_tuple:
+                arg_pos, rel_pos = mapping[j]
+                args_with_ind[arg_pos].indices[rel_pos] = i
+        self.args_with_ind: List[_ArgE] = args_with_ind
+        self.number_of_contraction_indices: int = len(contraction_indices)
+        self._track_permutation: Optional[List[List[int]]] = None
+
+        mapping = _get_mapping_from_subranks(base_array.subranks)
+
+        # Trick: add diagonalized indices as negative indices into the editor object:
+        for i, e in enumerate(diagonalized):
+            for j in e:
+                arg_pos, rel_pos = mapping[j]
+                self.args_with_ind[arg_pos].indices[rel_pos] = -1 - i
+
+    def insert_after(self, arg: _ArgE, new_arg: _ArgE):
+        pos = self.args_with_ind.index(arg)
+        self.args_with_ind.insert(pos + 1, new_arg)
+
+    def get_new_contraction_index(self):
+        self.number_of_contraction_indices += 1
+        return self.number_of_contraction_indices - 1
+
+    def refresh_indices(self):
+        updates = {}
+        for arg_with_ind in self.args_with_ind:
+            updates.update({i: -1 for i in arg_with_ind.indices if i is not None})
+        for i, e in enumerate(sorted(updates)):
+            updates[e] = i
+        self.number_of_contraction_indices = len(updates)
+        for arg_with_ind in self.args_with_ind:
+            arg_with_ind.indices = [updates.get(i, None) for i in arg_with_ind.indices]
+
+    def merge_scalars(self):
+        scalars = []
+        for arg_with_ind in self.args_with_ind:
+            if len(arg_with_ind.indices) == 0:
+                scalars.append(arg_with_ind)
+        for i in scalars:
+            self.args_with_ind.remove(i)
+        scalar = Mul.fromiter([i.element for i in scalars])
+        if len(self.args_with_ind) == 0:
+            self.args_with_ind.append(_ArgE(scalar))
+        else:
+            from sympy.tensor.array.expressions.from_array_to_matrix import _a2m_tensor_product
+            self.args_with_ind[0].element = _a2m_tensor_product(scalar, self.args_with_ind[0].element)
+
+    def to_array_contraction(self):
+
+        # Count the ranks of the arguments:
+        counter = 0
+        # Create a collector for the new diagonal indices:
+        diag_indices = defaultdict(list)
+
+        count_index_freq = Counter()
+        for arg_with_ind in self.args_with_ind:
+            count_index_freq.update(Counter(arg_with_ind.indices))
+
+        free_index_count = count_index_freq[None]
+
+        # Construct the inverse permutation:
+        inv_perm1 = []
+        inv_perm2 = []
+        # Keep track of which diagonal indices have already been processed:
+        done = set()
+
+        # Counter for the diagonal indices:
+        counter4 = 0
+
+        for arg_with_ind in self.args_with_ind:
+            # If some diagonalization axes have been removed, they should be
+            # permuted in order to keep the permutation.
+            # Add permutation here
+            counter2 = 0  # counter for the indices
+            for i in arg_with_ind.indices:
+                if i is None:
+                    inv_perm1.append(counter4)
+                    counter2 += 1
+                    counter4 += 1
+                    continue
+                if i >= 0:
+                    continue
+                # Reconstruct the diagonal indices:
+                diag_indices[-1 - i].append(counter + counter2)
+                if count_index_freq[i] == 1 and i not in done:
+                    inv_perm1.append(free_index_count - 1 - i)
+                    done.add(i)
+                elif i not in done:
+                    inv_perm2.append(free_index_count - 1 - i)
+                    done.add(i)
+                counter2 += 1
+            # Remove negative indices to restore a proper editor object:
+            arg_with_ind.indices = [i if i is not None and i >= 0 else None for i in arg_with_ind.indices]
+            counter += len([i for i in arg_with_ind.indices if i is None or i < 0])
+
+        inverse_permutation = inv_perm1 + inv_perm2
+        permutation = _af_invert(inverse_permutation)
+
+        # Get the diagonal indices after the detection of HadamardProduct in the expression:
+        diag_indices_filtered = [tuple(v) for v in diag_indices.values() if len(v) > 1]
+
+        self.merge_scalars()
+        self.refresh_indices()
+        args = [arg.element for arg in self.args_with_ind]
+        contraction_indices = self.get_contraction_indices()
+        expr = _array_contraction(_array_tensor_product(*args), *contraction_indices)
+        expr2 = _array_diagonal(expr, *diag_indices_filtered)
+        if self._track_permutation is not None:
+            permutation2 = _af_invert([j for i in self._track_permutation for j in i])
+            expr2 = _permute_dims(expr2, permutation2)
+
+        expr3 = _permute_dims(expr2, permutation)
+        return expr3
+
+    def get_contraction_indices(self) -> List[List[int]]:
+        contraction_indices: List[List[int]] = [[] for i in range(self.number_of_contraction_indices)]
+        current_position: int = 0
+        for arg_with_ind in self.args_with_ind:
+            for j in arg_with_ind.indices:
+                if j is not None:
+                    contraction_indices[j].append(current_position)
+                current_position += 1
+        return contraction_indices
+
+    def get_mapping_for_index(self, ind) -> List[_IndPos]:
+        if ind >= self.number_of_contraction_indices:
+            raise ValueError("index value exceeding the index range")
+        positions: List[_IndPos] = []
+        for i, arg_with_ind in enumerate(self.args_with_ind):
+            for j, arg_ind in enumerate(arg_with_ind.indices):
+                if ind == arg_ind:
+                    positions.append(_IndPos(i, j))
+        return positions
+
+    def get_contraction_indices_to_ind_rel_pos(self) -> List[List[_IndPos]]:
+        contraction_indices: List[List[_IndPos]] = [[] for i in range(self.number_of_contraction_indices)]
+        for i, arg_with_ind in enumerate(self.args_with_ind):
+            for j, ind in enumerate(arg_with_ind.indices):
+                if ind is not None:
+                    contraction_indices[ind].append(_IndPos(i, j))
+        return contraction_indices
+
+    def count_args_with_index(self, index: int) -> int:
+        """
+        Count the number of arguments that have the given index.
+        """
+        counter: int = 0
+        for arg_with_ind in self.args_with_ind:
+            if index in arg_with_ind.indices:
+                counter += 1
+        return counter
+
+    def get_args_with_index(self, index: int) -> List[_ArgE]:
+        """
+        Get a list of arguments having the given index.
+        """
+        ret: List[_ArgE] = [i for i in self.args_with_ind if index in i.indices]
+        return ret
+
+    @property
+    def number_of_diagonal_indices(self):
+        data = set()
+        for arg in self.args_with_ind:
+            data.update({i for i in arg.indices if i is not None and i < 0})
+        return len(data)
+
+    def track_permutation_start(self):
+        permutation = []
+        perm_diag = []
+        counter = 0
+        counter2 = -1
+        for arg_with_ind in self.args_with_ind:
+            perm = []
+            for i in arg_with_ind.indices:
+                if i is not None:
+                    if i < 0:
+                        perm_diag.append(counter2)
+                        counter2 -= 1
+                    continue
+                perm.append(counter)
+                counter += 1
+            permutation.append(perm)
+        max_ind = max(max(i) if i else -1 for i in permutation) if permutation else -1
+        perm_diag = [max_ind - i for i in perm_diag]
+        self._track_permutation = permutation + [perm_diag]
+
+    def track_permutation_merge(self, destination: _ArgE, from_element: _ArgE):
+        index_destination = self.args_with_ind.index(destination)
+        index_element = self.args_with_ind.index(from_element)
+        self._track_permutation[index_destination].extend(self._track_permutation[index_element]) # type: ignore
+        self._track_permutation.pop(index_element) # type: ignore
+
+    def get_absolute_free_range(self, arg: _ArgE) -> typing.Tuple[int, int]:
+        """
+        Return the range of the free indices of the arg as absolute positions
+        among all free indices.
+        """
+        counter = 0
+        for arg_with_ind in self.args_with_ind:
+            number_free_indices = len([i for i in arg_with_ind.indices if i is None])
+            if arg_with_ind == arg:
+                return counter, counter + number_free_indices
+            counter += number_free_indices
+        raise IndexError("argument not found")
+
+    def get_absolute_range(self, arg: _ArgE) -> typing.Tuple[int, int]:
+        """
+        Return the absolute range of indices for arg, disregarding dummy
+        indices.
+        """
+        counter = 0
+        for arg_with_ind in self.args_with_ind:
+            number_indices = len(arg_with_ind.indices)
+            if arg_with_ind == arg:
+                return counter, counter + number_indices
+            counter += number_indices
+        raise IndexError("argument not found")
+
+
+def get_rank(expr):
+    if isinstance(expr, (MatrixExpr, MatrixElement)):
+        return 2
+    if isinstance(expr, _CodegenArrayAbstract):
+        return len(expr.shape)
+    if isinstance(expr, NDimArray):
+        return expr.rank()
+    if isinstance(expr, Indexed):
+        return expr.rank
+    if isinstance(expr, IndexedBase):
+        shape = expr.shape
+        if shape is None:
+            return -1
+        else:
+            return len(shape)
+    if hasattr(expr, "shape"):
+        return len(expr.shape)
+    return 0
+
+
+def _get_subrank(expr):
+    if isinstance(expr, _CodegenArrayAbstract):
+        return expr.subrank()
+    return get_rank(expr)
+
+
+def _get_subranks(expr):
+    if isinstance(expr, _CodegenArrayAbstract):
+        return expr.subranks
+    else:
+        return [get_rank(expr)]
+
+
+def get_shape(expr):
+    if hasattr(expr, "shape"):
+        return expr.shape
+    return ()
+
+
+def nest_permutation(expr):
+    if isinstance(expr, PermuteDims):
+        return expr.nest_permutation()
+    else:
+        return expr
+
+
+def _array_tensor_product(*args, **kwargs):
+    return ArrayTensorProduct(*args, canonicalize=True, **kwargs)
+
+
+def _array_contraction(expr, *contraction_indices, **kwargs):
+    return ArrayContraction(expr, *contraction_indices, canonicalize=True, **kwargs)
+
+
+def _array_diagonal(expr, *diagonal_indices, **kwargs):
+    return ArrayDiagonal(expr, *diagonal_indices, canonicalize=True, **kwargs)
+
+
+def _permute_dims(expr, permutation, **kwargs):
+    return PermuteDims(expr, permutation, canonicalize=True, **kwargs)
+
+
+def _array_add(*args, **kwargs):
+    return ArrayAdd(*args, canonicalize=True, **kwargs)
+
+
+def _get_array_element_or_slice(expr, indices):
+    return ArrayElement(expr, indices)

.venv/lib/python3.11/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()

.venv/lib/python3.11/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)

.venv/lib/python3.11/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)

.venv/lib/python3.11/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)

.venv/lib/python3.11/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)

.venv/lib/python3.11/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)

.venv/lib/python3.11/site-packages/sympy/tensor/array/expressions/from_array_to_matrix.py

@@ -0,0 +1,1003 @@
+import itertools
+from collections import defaultdict
+from typing import Tuple as tTuple, Union as tUnion, FrozenSet, Dict as tDict, List, Optional
+from functools import singledispatch
+from itertools import accumulate
+
+from sympy import MatMul, Basic, Wild, KroneckerProduct
+from sympy.assumptions.ask import (Q, ask)
+from sympy.core.mul import Mul
+from sympy.core.singleton import S
+from sympy.matrices.expressions.diagonal import DiagMatrix
+from sympy.matrices.expressions.hadamard import hadamard_product, HadamardPower
+from sympy.matrices.expressions.matexpr import MatrixExpr
+from sympy.matrices.expressions.special import (Identity, ZeroMatrix, OneMatrix)
+from sympy.matrices.expressions.trace import Trace
+from sympy.matrices.expressions.transpose import Transpose
+from sympy.combinatorics.permutations import _af_invert, Permutation
+from sympy.matrices.matrixbase import MatrixBase
+from sympy.matrices.expressions.applyfunc import ElementwiseApplyFunction
+from sympy.matrices.expressions.matexpr import MatrixElement
+from sympy.tensor.array.expressions.array_expressions import PermuteDims, ArrayDiagonal, \
+    ArrayTensorProduct, OneArray, get_rank, _get_subrank, ZeroArray, ArrayContraction, \
+    ArrayAdd, _CodegenArrayAbstract, get_shape, ArrayElementwiseApplyFunc, _ArrayExpr, _EditArrayContraction, _ArgE, \
+    ArrayElement, _array_tensor_product, _array_contraction, _array_diagonal, _array_add, _permute_dims
+from sympy.tensor.array.expressions.utils import _get_mapping_from_subranks
+
+
+def _get_candidate_for_matmul_from_contraction(scan_indices: List[Optional[int]], remaining_args: List[_ArgE]) -> tTuple[Optional[_ArgE], bool, int]:
+
+    scan_indices_int: List[int] = [i for i in scan_indices if i is not None]
+    if len(scan_indices_int) == 0:
+        return None, False, -1
+
+    transpose: bool = False
+    candidate: Optional[_ArgE] = None
+    candidate_index: int = -1
+    for arg_with_ind2 in remaining_args:
+        if not isinstance(arg_with_ind2.element, MatrixExpr):
+            continue
+        for index in scan_indices_int:
+            if candidate_index != -1 and candidate_index != index:
+                # A candidate index has already been selected, check
+                # repetitions only for that index:
+                continue
+            if index in arg_with_ind2.indices:
+                if set(arg_with_ind2.indices) == {index}:
+                    # Index repeated twice in arg_with_ind2
+                    candidate = None
+                    break
+                if candidate is None:
+                    candidate = arg_with_ind2
+                    candidate_index = index
+                    transpose = (index == arg_with_ind2.indices[1])
+                else:
+                    # Index repeated more than twice, break
+                    candidate = None
+                    break
+    return candidate, transpose, candidate_index
+
+
+def _insert_candidate_into_editor(editor: _EditArrayContraction, arg_with_ind: _ArgE, candidate: _ArgE, transpose1: bool, transpose2: bool):
+    other = candidate.element
+    other_index: Optional[int]
+    if transpose2:
+        other = Transpose(other)
+        other_index = candidate.indices[0]
+    else:
+        other_index = candidate.indices[1]
+    new_element = (Transpose(arg_with_ind.element) if transpose1 else arg_with_ind.element) * other
+    editor.args_with_ind.remove(candidate)
+    new_arge = _ArgE(new_element)
+    return new_arge, other_index
+
+
+def _support_function_tp1_recognize(contraction_indices, args):
+    if len(contraction_indices) == 0:
+        return _a2m_tensor_product(*args)
+
+    ac = _array_contraction(_array_tensor_product(*args), *contraction_indices)
+    editor = _EditArrayContraction(ac)
+    editor.track_permutation_start()
+
+    while True:
+        flag_stop = True
+        for i, arg_with_ind in enumerate(editor.args_with_ind):
+            if not isinstance(arg_with_ind.element, MatrixExpr):
+                continue
+
+            first_index = arg_with_ind.indices[0]
+            second_index = arg_with_ind.indices[1]
+
+            first_frequency = editor.count_args_with_index(first_index)
+            second_frequency = editor.count_args_with_index(second_index)
+
+            if first_index is not None and first_frequency == 1 and first_index == second_index:
+                flag_stop = False
+                arg_with_ind.element = Trace(arg_with_ind.element)._normalize()
+                arg_with_ind.indices = []
+                break
+
+            scan_indices = []
+            if first_frequency == 2:
+                scan_indices.append(first_index)
+            if second_frequency == 2:
+                scan_indices.append(second_index)
+
+            candidate, transpose, found_index = _get_candidate_for_matmul_from_contraction(scan_indices, editor.args_with_ind[i+1:])
+            if candidate is not None:
+                flag_stop = False
+                editor.track_permutation_merge(arg_with_ind, candidate)
+                transpose1 = found_index == first_index
+                new_arge, other_index = _insert_candidate_into_editor(editor, arg_with_ind, candidate, transpose1, transpose)
+                if found_index == first_index:
+                    new_arge.indices = [second_index, other_index]
+                else:
+                    new_arge.indices = [first_index, other_index]
+                set_indices = set(new_arge.indices)
+                if len(set_indices) == 1 and set_indices != {None}:
+                    # This is a trace:
+                    new_arge.element = Trace(new_arge.element)._normalize()
+                    new_arge.indices = []
+                editor.args_with_ind[i] = new_arge
+                # TODO: is this break necessary?
+                break
+
+        if flag_stop:
+            break
+
+    editor.refresh_indices()
+    return editor.to_array_contraction()
+
+
+def _find_trivial_matrices_rewrite(expr: ArrayTensorProduct):
+    # If there are matrices of trivial shape in the tensor product (i.e. shape
+    # (1, 1)), try to check if there is a suitable non-trivial MatMul where the
+    # expression can be inserted.
+
+    # For example, if "a" has shape (1, 1) and "b" has shape (k, 1), the
+    # expressions "_array_tensor_product(a, b*b.T)" can be rewritten as
+    # "b*a*b.T"
+
+    trivial_matrices = []
+    pos: Optional[int] = None
+    first: Optional[MatrixExpr] = None
+    second: Optional[MatrixExpr] = None
+    removed: List[int] = []
+    counter: int = 0
+    args: List[Optional[Basic]] = list(expr.args)
+    for i, arg in enumerate(expr.args):
+        if isinstance(arg, MatrixExpr):
+            if arg.shape == (1, 1):
+                trivial_matrices.append(arg)
+                args[i] = None
+                removed.extend([counter, counter+1])
+            elif pos is None and isinstance(arg, MatMul):
+                margs = arg.args
+                for j, e in enumerate(margs):
+                    if isinstance(e, MatrixExpr) and e.shape[1] == 1:
+                        pos = i
+                        first = MatMul.fromiter(margs[:j+1])
+                        second = MatMul.fromiter(margs[j+1:])
+                        break
+        counter += get_rank(arg)
+    if pos is None:
+        return expr, []
+    args[pos] = (first*MatMul.fromiter(i for i in trivial_matrices)*second).doit()
+    return _array_tensor_product(*[i for i in args if i is not None]), removed
+
+
+def _find_trivial_kronecker_products_broadcast(expr: ArrayTensorProduct):
+    newargs: List[Basic] = []
+    removed = []
+    count_dims = 0
+    for arg in expr.args:
+        count_dims += get_rank(arg)
+        shape = get_shape(arg)
+        current_range = [count_dims-i for i in range(len(shape), 0, -1)]
+        if (shape == (1, 1) and len(newargs) > 0 and 1 not in get_shape(newargs[-1]) and
+            isinstance(newargs[-1], MatrixExpr) and isinstance(arg, MatrixExpr)):
+            # KroneckerProduct object allows the trick of broadcasting:
+            newargs[-1] = KroneckerProduct(newargs[-1], arg)
+            removed.extend(current_range)
+        elif 1 not in shape and len(newargs) > 0 and get_shape(newargs[-1]) == (1, 1):
+            # Broadcast:
+            newargs[-1] = KroneckerProduct(newargs[-1], arg)
+            prev_range = [i for i in range(min(current_range)) if i not in removed]
+            removed.extend(prev_range[-2:])
+        else:
+            newargs.append(arg)
+    return _array_tensor_product(*newargs), removed
+
+
+@singledispatch
+def _array2matrix(expr):
+    return expr
+
+
+@_array2matrix.register(ZeroArray)
+def _(expr: ZeroArray):
+    if get_rank(expr) == 2:
+        return ZeroMatrix(*expr.shape)
+    else:
+        return expr
+
+
+@_array2matrix.register(ArrayTensorProduct)
+def _(expr: ArrayTensorProduct):
+    return _a2m_tensor_product(*[_array2matrix(arg) for arg in expr.args])
+
+
+@_array2matrix.register(ArrayContraction)
+def _(expr: ArrayContraction):
+    expr = expr.flatten_contraction_of_diagonal()
+    expr = identify_removable_identity_matrices(expr)
+    expr = expr.split_multiple_contractions()
+    expr = identify_hadamard_products(expr)
+    if not isinstance(expr, ArrayContraction):
+        return _array2matrix(expr)
+    subexpr = expr.expr
+    contraction_indices: tTuple[tTuple[int]] = expr.contraction_indices
+    if contraction_indices == ((0,), (1,)) or (
+        contraction_indices == ((0,),) and subexpr.shape[1] == 1
+    ) or (
+        contraction_indices == ((1,),) and subexpr.shape[0] == 1
+    ):
+        shape = subexpr.shape
+        subexpr = _array2matrix(subexpr)
+        if isinstance(subexpr, MatrixExpr):
+            return OneMatrix(1, shape[0])*subexpr*OneMatrix(shape[1], 1)
+    if isinstance(subexpr, ArrayTensorProduct):
+        newexpr = _array_contraction(_array2matrix(subexpr), *contraction_indices)
+        contraction_indices = newexpr.contraction_indices
+        if any(i > 2 for i in newexpr.subranks):
+            addends = _array_add(*[_a2m_tensor_product(*j) for j in itertools.product(*[i.args if isinstance(i,
+                                                                                                                             ArrayAdd) else [i] for i in expr.expr.args])])
+            newexpr = _array_contraction(addends, *contraction_indices)
+        if isinstance(newexpr, ArrayAdd):
+            ret = _array2matrix(newexpr)
+            return ret
+        assert isinstance(newexpr, ArrayContraction)
+        ret = _support_function_tp1_recognize(contraction_indices, list(newexpr.expr.args))
+        return ret
+    elif not isinstance(subexpr, _CodegenArrayAbstract):
+        ret = _array2matrix(subexpr)
+        if isinstance(ret, MatrixExpr):
+            assert expr.contraction_indices == ((0, 1),)
+            return _a2m_trace(ret)
+        else:
+            return _array_contraction(ret, *expr.contraction_indices)
+
+
+@_array2matrix.register(ArrayDiagonal)
+def _(expr: ArrayDiagonal):
+    pexpr = _array_diagonal(_array2matrix(expr.expr), *expr.diagonal_indices)
+    pexpr = identify_hadamard_products(pexpr)
+    if isinstance(pexpr, ArrayDiagonal):
+        pexpr = _array_diag2contr_diagmatrix(pexpr)
+    if expr == pexpr:
+        return expr
+    return _array2matrix(pexpr)
+
+
+@_array2matrix.register(PermuteDims)
+def _(expr: PermuteDims):
+    if expr.permutation.array_form == [1, 0]:
+        return _a2m_transpose(_array2matrix(expr.expr))
+    elif isinstance(expr.expr, ArrayTensorProduct):
+        ranks = expr.expr.subranks
+        inv_permutation = expr.permutation**(-1)
+        newrange = [inv_permutation(i) for i in range(sum(ranks))]
+        newpos = []
+        counter = 0
+        for rank in ranks:
+            newpos.append(newrange[counter:counter+rank])
+            counter += rank
+        newargs = []
+        newperm = []
+        scalars = []
+        for pos, arg in zip(newpos, expr.expr.args):
+            if len(pos) == 0:
+                scalars.append(_array2matrix(arg))
+            elif pos == sorted(pos):
+                newargs.append((_array2matrix(arg), pos[0]))
+                newperm.extend(pos)
+            elif len(pos) == 2:
+                newargs.append((_a2m_transpose(_array2matrix(arg)), pos[0]))
+                newperm.extend(reversed(pos))
+            else:
+                raise NotImplementedError()
+        newargs = [i[0] for i in newargs]
+        return _permute_dims(_a2m_tensor_product(*scalars, *newargs), _af_invert(newperm))
+    elif isinstance(expr.expr, ArrayContraction):
+        mat_mul_lines = _array2matrix(expr.expr)
+        if not isinstance(mat_mul_lines, ArrayTensorProduct):
+            return _permute_dims(mat_mul_lines, expr.permutation)
+        # TODO: this assumes that all arguments are matrices, it may not be the case:
+        permutation = Permutation(2*len(mat_mul_lines.args)-1)*expr.permutation
+        permuted = [permutation(i) for i in range(2*len(mat_mul_lines.args))]
+        args_array = [None for i in mat_mul_lines.args]
+        for i in range(len(mat_mul_lines.args)):
+            p1 = permuted[2*i]
+            p2 = permuted[2*i+1]
+            if p1 // 2 != p2 // 2:
+                return _permute_dims(mat_mul_lines, permutation)
+            if p1 > p2:
+                args_array[i] = _a2m_transpose(mat_mul_lines.args[p1 // 2])
+            else:
+                args_array[i] = mat_mul_lines.args[p1 // 2]
+        return _a2m_tensor_product(*args_array)
+    else:
+        return expr
+
+
+@_array2matrix.register(ArrayAdd)
+def _(expr: ArrayAdd):
+    addends = [_array2matrix(arg) for arg in expr.args]
+    return _a2m_add(*addends)
+
+
+@_array2matrix.register(ArrayElementwiseApplyFunc)
+def _(expr: ArrayElementwiseApplyFunc):
+    subexpr = _array2matrix(expr.expr)
+    if isinstance(subexpr, MatrixExpr):
+        if subexpr.shape != (1, 1):
+            d = expr.function.bound_symbols[0]
+            w = Wild("w", exclude=[d])
+            p = Wild("p", exclude=[d])
+            m = expr.function.expr.match(w*d**p)
+            if m is not None:
+                return m[w]*HadamardPower(subexpr, m[p])
+        return ElementwiseApplyFunction(expr.function, subexpr)
+    else:
+        return ArrayElementwiseApplyFunc(expr.function, subexpr)
+
+
+@_array2matrix.register(ArrayElement)
+def _(expr: ArrayElement):
+    ret = _array2matrix(expr.name)
+    if isinstance(ret, MatrixExpr):
+        return MatrixElement(ret, *expr.indices)
+    return ArrayElement(ret, expr.indices)
+
+
+@singledispatch
+def _remove_trivial_dims(expr):
+    return expr, []
+
+
+@_remove_trivial_dims.register(ArrayTensorProduct)
+def _(expr: ArrayTensorProduct):
+    # Recognize expressions like [x, y] with shape (k, 1, k, 1) as `x*y.T`.
+    # The matrix expression has to be equivalent to the tensor product of the
+    # matrices, with trivial dimensions (i.e. dim=1) dropped.
+    # That is, add contractions over trivial dimensions:
+
+    removed = []
+    newargs = []
+    cumul = list(accumulate([0] + [get_rank(arg) for arg in expr.args]))
+    pending = None
+    prev_i = None
+    for i, arg in enumerate(expr.args):
+        current_range = list(range(cumul[i], cumul[i+1]))
+        if isinstance(arg, OneArray):
+            removed.extend(current_range)
+            continue
+        if not isinstance(arg, (MatrixExpr, MatrixBase)):
+            rarg, rem = _remove_trivial_dims(arg)
+            removed.extend(rem)
+            newargs.append(rarg)
+            continue
+        elif getattr(arg, "is_Identity", False) and arg.shape == (1, 1):
+            if arg.shape == (1, 1):
+                # Ignore identity matrices of shape (1, 1) - they are equivalent to scalar 1.
+                removed.extend(current_range)
+            continue
+        elif arg.shape == (1, 1):
+            arg, _ = _remove_trivial_dims(arg)
+            # Matrix is equivalent to scalar:
+            if len(newargs) == 0:
+                newargs.append(arg)
+            elif 1 in get_shape(newargs[-1]):
+                if newargs[-1].shape[1] == 1:
+                    newargs[-1] = newargs[-1]*arg
+                else:
+                    newargs[-1] = arg*newargs[-1]
+                removed.extend(current_range)
+            else:
+                newargs.append(arg)
+        elif 1 in arg.shape:
+            k = [i for i in arg.shape if i != 1][0]
+            if pending is None:
+                pending = k
+                prev_i = i
+                newargs.append(arg)
+            elif pending == k:
+                prev = newargs[-1]
+                if prev.shape[0] == 1:
+                    d1 = cumul[prev_i]
+                    prev = _a2m_transpose(prev)
+                else:
+                    d1 = cumul[prev_i] + 1
+                if arg.shape[1] == 1:
+                    d2 = cumul[i] + 1
+                    arg = _a2m_transpose(arg)
+                else:
+                    d2 = cumul[i]
+                newargs[-1] = prev*arg
+                pending = None
+                removed.extend([d1, d2])
+            else:
+                newargs.append(arg)
+                pending = k
+                prev_i = i
+        else:
+            newargs.append(arg)
+            pending = None
+    newexpr, newremoved = _a2m_tensor_product(*newargs), sorted(removed)
+    if isinstance(newexpr, ArrayTensorProduct):
+        newexpr, newremoved2 = _find_trivial_matrices_rewrite(newexpr)
+        newremoved = _combine_removed(-1, newremoved, newremoved2)
+    if isinstance(newexpr, ArrayTensorProduct):
+        newexpr, newremoved2 = _find_trivial_kronecker_products_broadcast(newexpr)
+        newremoved = _combine_removed(-1, newremoved, newremoved2)
+    return newexpr, newremoved
+
+
+@_remove_trivial_dims.register(ArrayAdd)
+def _(expr: ArrayAdd):
+    rec = [_remove_trivial_dims(arg) for arg in expr.args]
+    newargs, removed = zip(*rec)
+    if len({get_shape(i) for i in newargs}) > 1:
+        return expr, []
+    if len(removed) == 0:
+        return expr, removed
+    removed1 = removed[0]
+    return _a2m_add(*newargs), removed1
+
+
+@_remove_trivial_dims.register(PermuteDims)
+def _(expr: PermuteDims):
+    subexpr, subremoved = _remove_trivial_dims(expr.expr)
+    p = expr.permutation.array_form
+    pinv = _af_invert(expr.permutation.array_form)
+    shift = list(accumulate([1 if i in subremoved else 0 for i in range(len(p))]))
+    premoved = [pinv[i] for i in subremoved]
+    p2 = [e - shift[e] for e in p if e not in subremoved]
+    # TODO: check if subremoved should be permuted as well...
+    newexpr = _permute_dims(subexpr, p2)
+    premoved = sorted(premoved)
+    if newexpr != expr:
+        newexpr, removed2 = _remove_trivial_dims(_array2matrix(newexpr))
+        premoved = _combine_removed(-1, premoved, removed2)
+    return newexpr, premoved
+
+
+@_remove_trivial_dims.register(ArrayContraction)
+def _(expr: ArrayContraction):
+    new_expr, removed0 = _array_contraction_to_diagonal_multiple_identity(expr)
+    if new_expr != expr:
+        new_expr2, removed1 = _remove_trivial_dims(_array2matrix(new_expr))
+        removed = _combine_removed(-1, removed0, removed1)
+        return new_expr2, removed
+    rank1 = get_rank(expr)
+    expr, removed1 = remove_identity_matrices(expr)
+    if not isinstance(expr, ArrayContraction):
+        expr2, removed2 = _remove_trivial_dims(expr)
+        return expr2, _combine_removed(rank1, removed1, removed2)
+    newexpr, removed2 = _remove_trivial_dims(expr.expr)
+    shifts = list(accumulate([1 if i in removed2 else 0 for i in range(get_rank(expr.expr))]))
+    new_contraction_indices = [tuple(j for j in i if j not in removed2) for i in expr.contraction_indices]
+    # Remove possible empty tuples "()":
+    new_contraction_indices = [i for i in new_contraction_indices if len(i) > 0]
+    contraction_indices_flat = [j for i in expr.contraction_indices for j in i]
+    removed2 = [i for i in removed2 if i not in contraction_indices_flat]
+    new_contraction_indices = [tuple(j - shifts[j] for j in i) for i in new_contraction_indices]
+    # Shift removed2:
+    removed2 = ArrayContraction._push_indices_up(expr.contraction_indices, removed2)
+    removed = _combine_removed(rank1, removed1, removed2)
+    return _array_contraction(newexpr, *new_contraction_indices), list(removed)
+
+
+def _remove_diagonalized_identity_matrices(expr: ArrayDiagonal):
+    assert isinstance(expr, ArrayDiagonal)
+    editor = _EditArrayContraction(expr)
+    mapping = {i: {j for j in editor.args_with_ind if i in j.indices} for i in range(-1, -1-editor.number_of_diagonal_indices, -1)}
+    removed = []
+    counter: int = 0
+    for i, arg_with_ind in enumerate(editor.args_with_ind):
+        counter += len(arg_with_ind.indices)
+        if isinstance(arg_with_ind.element, Identity):
+            if None in arg_with_ind.indices and any(i is not None and (i < 0) == True for i in arg_with_ind.indices):
+                diag_ind = [j for j in arg_with_ind.indices if j is not None][0]
+                other = [j for j in mapping[diag_ind] if j != arg_with_ind][0]
+                if not isinstance(other.element, MatrixExpr):
+                    continue
+                if 1 not in other.element.shape:
+                    continue
+                if None not in other.indices:
+                    continue
+                editor.args_with_ind[i].element = None
+                none_index = other.indices.index(None)
+                other.element = DiagMatrix(other.element)
+                other_range = editor.get_absolute_range(other)
+                removed.extend([other_range[0] + none_index])
+    editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None]
+    removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, get_rank(expr.expr))
+    return editor.to_array_contraction(), removed
+
+
+@_remove_trivial_dims.register(ArrayDiagonal)
+def _(expr: ArrayDiagonal):
+    newexpr, removed = _remove_trivial_dims(expr.expr)
+    shifts = list(accumulate([0] + [1 if i in removed else 0 for i in range(get_rank(expr.expr))]))
+    new_diag_indices_map = {i: tuple(j for j in i if j not in removed) for i in expr.diagonal_indices}
+    for old_diag_tuple, new_diag_tuple in new_diag_indices_map.items():
+        if len(new_diag_tuple) == 1:
+            removed = [i for i in removed if i not in old_diag_tuple]
+    new_diag_indices = [tuple(j - shifts[j] for j in i) for i in new_diag_indices_map.values()]
+    rank = get_rank(expr.expr)
+    removed = ArrayDiagonal._push_indices_up(expr.diagonal_indices, removed, rank)
+    removed = sorted(set(removed))
+    # If there are single axes to diagonalize remaining, it means that their
+    # corresponding dimension has been removed, they no longer need diagonalization:
+    new_diag_indices = [i for i in new_diag_indices if len(i) > 0]
+    if len(new_diag_indices) > 0:
+        newexpr2 = _array_diagonal(newexpr, *new_diag_indices, allow_trivial_diags=True)
+    else:
+        newexpr2 = newexpr
+    if isinstance(newexpr2, ArrayDiagonal):
+        newexpr3, removed2 = _remove_diagonalized_identity_matrices(newexpr2)
+        removed = _combine_removed(-1, removed, removed2)
+        return newexpr3, removed
+    else:
+        return newexpr2, removed
+
+
+@_remove_trivial_dims.register(ElementwiseApplyFunction)
+def _(expr: ElementwiseApplyFunction):
+    subexpr, removed = _remove_trivial_dims(expr.expr)
+    if subexpr.shape == (1, 1):
+        # TODO: move this to ElementwiseApplyFunction
+        return expr.function(subexpr), removed + [0, 1]
+    return ElementwiseApplyFunction(expr.function, subexpr), []
+
+
+@_remove_trivial_dims.register(ArrayElementwiseApplyFunc)
+def _(expr: ArrayElementwiseApplyFunc):
+    subexpr, removed = _remove_trivial_dims(expr.expr)
+    return ArrayElementwiseApplyFunc(expr.function, subexpr), removed
+
+
+def convert_array_to_matrix(expr):
+    r"""
+    Recognize matrix expressions in codegen objects.
+
+    If more than one matrix multiplication line have been detected, return a
+    list with the matrix expressions.
+
+    Examples
+    ========
+
+    >>> from sympy.tensor.array.expressions.from_indexed_to_array import convert_indexed_to_array
+    >>> from sympy.tensor.array import tensorcontraction, tensorproduct
+    >>> from sympy import MatrixSymbol, Sum
+    >>> from sympy.abc import i, j, k, l, N
+    >>> from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+    >>> from sympy.tensor.array.expressions.from_array_to_matrix import convert_array_to_matrix
+    >>> A = MatrixSymbol("A", N, N)
+    >>> B = MatrixSymbol("B", N, N)
+    >>> C = MatrixSymbol("C", N, N)
+    >>> D = MatrixSymbol("D", N, N)
+
+    >>> expr = Sum(A[i, j]*B[j, k], (j, 0, N-1))
+    >>> cg = convert_indexed_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    A*B
+    >>> cg = convert_indexed_to_array(expr, first_indices=[k])
+    >>> convert_array_to_matrix(cg)
+    B.T*A.T
+
+    Transposition is detected:
+
+    >>> expr = Sum(A[j, i]*B[j, k], (j, 0, N-1))
+    >>> cg = convert_indexed_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    A.T*B
+    >>> cg = convert_indexed_to_array(expr, first_indices=[k])
+    >>> convert_array_to_matrix(cg)
+    B.T*A
+
+    Detect the trace:
+
+    >>> expr = Sum(A[i, i], (i, 0, N-1))
+    >>> cg = convert_indexed_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    Trace(A)
+
+    Recognize some more complex traces:
+
+    >>> expr = Sum(A[i, j]*B[j, i], (i, 0, N-1), (j, 0, N-1))
+    >>> cg = convert_indexed_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    Trace(A*B)
+
+    More complicated expressions:
+
+    >>> expr = Sum(A[i, j]*B[k, j]*A[l, k], (j, 0, N-1), (k, 0, N-1))
+    >>> cg = convert_indexed_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    A*B.T*A.T
+
+    Expressions constructed from matrix expressions do not contain literal
+    indices, the positions of free indices are returned instead:
+
+    >>> expr = A*B
+    >>> cg = convert_matrix_to_array(expr)
+    >>> convert_array_to_matrix(cg)
+    A*B
+
+    If more than one line of matrix multiplications is detected, return
+    separate matrix multiplication factors embedded in a tensor product object:
+
+    >>> cg = tensorcontraction(tensorproduct(A, B, C, D), (1, 2), (5, 6))
+    >>> convert_array_to_matrix(cg)
+    ArrayTensorProduct(A*B, C*D)
+
+    The two lines have free indices at axes 0, 3 and 4, 7, respectively.
+    """
+    rec = _array2matrix(expr)
+    rec, removed = _remove_trivial_dims(rec)
+    return rec
+
+
+def _array_diag2contr_diagmatrix(expr: ArrayDiagonal):
+    if isinstance(expr.expr, ArrayTensorProduct):
+        args = list(expr.expr.args)
+        diag_indices = list(expr.diagonal_indices)
+        mapping = _get_mapping_from_subranks([_get_subrank(arg) for arg in args])
+        tuple_links = [[mapping[j] for j in i] for i in diag_indices]
+        contr_indices = []
+        total_rank = get_rank(expr)
+        replaced = [False for arg in args]
+        for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)):
+            if len(abs_pos) != 2:
+                continue
+            (pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos
+            arg1 = args[pos1_outer]
+            arg2 = args[pos2_outer]
+            if get_rank(arg1) != 2 or get_rank(arg2) != 2:
+                if replaced[pos1_outer]:
+                    diag_indices[i] = None
+                if replaced[pos2_outer]:
+                    diag_indices[i] = None
+                continue
+            pos1_in2 = 1 - pos1_inner
+            pos2_in2 = 1 - pos2_inner
+            if arg1.shape[pos1_in2] == 1:
+                if arg1.shape[pos1_inner] != 1:
+                    darg1 = DiagMatrix(arg1)
+                else:
+                    darg1 = arg1
+                args.append(darg1)
+                contr_indices.append(((pos2_outer, pos2_inner), (len(args)-1, pos1_inner)))
+                total_rank += 1
+                diag_indices[i] = None
+                args[pos1_outer] = OneArray(arg1.shape[pos1_in2])
+                replaced[pos1_outer] = True
+            elif arg2.shape[pos2_in2] == 1:
+                if arg2.shape[pos2_inner] != 1:
+                    darg2 = DiagMatrix(arg2)
+                else:
+                    darg2 = arg2
+                args.append(darg2)
+                contr_indices.append(((pos1_outer, pos1_inner), (len(args)-1, pos2_inner)))
+                total_rank += 1
+                diag_indices[i] = None
+                args[pos2_outer] = OneArray(arg2.shape[pos2_in2])
+                replaced[pos2_outer] = True
+        diag_indices_new = [i for i in diag_indices if i is not None]
+        cumul = list(accumulate([0] + [get_rank(arg) for arg in args]))
+        contr_indices2 = [tuple(cumul[a] + b for a, b in i) for i in contr_indices]
+        tc = _array_contraction(
+            _array_tensor_product(*args), *contr_indices2
+        )
+        td = _array_diagonal(tc, *diag_indices_new)
+        return td
+    return expr
+
+
+def _a2m_mul(*args):
+    if not any(isinstance(i, _CodegenArrayAbstract) for i in args):
+        from sympy.matrices.expressions.matmul import MatMul
+        return MatMul(*args).doit()
+    else:
+        return _array_contraction(
+            _array_tensor_product(*args),
+            *[(2*i-1, 2*i) for i in range(1, len(args))]
+        )
+
+
+def _a2m_tensor_product(*args):
+    scalars = []
+    arrays = []
+    for arg in args:
+        if isinstance(arg, (MatrixExpr, _ArrayExpr, _CodegenArrayAbstract)):
+            arrays.append(arg)
+        else:
+            scalars.append(arg)
+    scalar = Mul.fromiter(scalars)
+    if len(arrays) == 0:
+        return scalar
+    if scalar != 1:
+        if isinstance(arrays[0], _CodegenArrayAbstract):
+            arrays = [scalar] + arrays
+        else:
+            arrays[0] *= scalar
+    return _array_tensor_product(*arrays)
+
+
+def _a2m_add(*args):
+    if not any(isinstance(i, _CodegenArrayAbstract) for i in args):
+        from sympy.matrices.expressions.matadd import MatAdd
+        return MatAdd(*args).doit()
+    else:
+        return _array_add(*args)
+
+
+def _a2m_trace(arg):
+    if isinstance(arg, _CodegenArrayAbstract):
+        return _array_contraction(arg, (0, 1))
+    else:
+        from sympy.matrices.expressions.trace import Trace
+        return Trace(arg)
+
+
+def _a2m_transpose(arg):
+    if isinstance(arg, _CodegenArrayAbstract):
+        return _permute_dims(arg, [1, 0])
+    else:
+        from sympy.matrices.expressions.transpose import Transpose
+        return Transpose(arg).doit()
+
+
+def identify_hadamard_products(expr: tUnion[ArrayContraction, ArrayDiagonal]):
+
+    editor: _EditArrayContraction = _EditArrayContraction(expr)
+
+    map_contr_to_args: tDict[FrozenSet, List[_ArgE]] = defaultdict(list)
+    map_ind_to_inds: tDict[Optional[int], int] = defaultdict(int)
+    for arg_with_ind in editor.args_with_ind:
+        for ind in arg_with_ind.indices:
+            map_ind_to_inds[ind] += 1
+        if None in arg_with_ind.indices:
+            continue
+        map_contr_to_args[frozenset(arg_with_ind.indices)].append(arg_with_ind)
+
+    k: FrozenSet[int]
+    v: List[_ArgE]
+    for k, v in map_contr_to_args.items():
+        make_trace: bool = False
+        if len(k) == 1 and next(iter(k)) >= 0 and sum(next(iter(k)) in i for i in map_contr_to_args) == 1:
+            # This is a trace: the arguments are fully contracted with only one
+            # index, and the index isn't used anywhere else:
+            make_trace = True
+            first_element = S.One
+        elif len(k) != 2:
+            # Hadamard product only defined for matrices:
+            continue
+        if len(v) == 1:
+            # Hadamard product with a single argument makes no sense:
+            continue
+        for ind in k:
+            if map_ind_to_inds[ind] <= 2:
+                # There is no other contraction, skip:
+                continue
+
+        def check_transpose(x):
+            x = [i if i >= 0 else -1-i for i in x]
+            return x == sorted(x)
+
+        # Check if expression is a trace:
+        if all(map_ind_to_inds[j] == len(v) and j >= 0 for j in k) and all(j >= 0 for j in k):
+            # This is a trace
+            make_trace = True
+            first_element = v[0].element
+            if not check_transpose(v[0].indices):
+                first_element = first_element.T
+            hadamard_factors = v[1:]
+        else:
+            hadamard_factors = v
+
+        # This is a Hadamard product:
+
+        hp = hadamard_product(*[i.element if check_transpose(i.indices) else Transpose(i.element) for i in hadamard_factors])
+        hp_indices = v[0].indices
+        if not check_transpose(hadamard_factors[0].indices):
+            hp_indices = list(reversed(hp_indices))
+        if make_trace:
+            hp = Trace(first_element*hp.T)._normalize()
+            hp_indices = []
+        editor.insert_after(v[0], _ArgE(hp, hp_indices))
+        for i in v:
+            editor.args_with_ind.remove(i)
+
+    return editor.to_array_contraction()
+
+
+def identify_removable_identity_matrices(expr):
+    editor = _EditArrayContraction(expr)
+
+    flag = True
+    while flag:
+        flag = False
+        for arg_with_ind in editor.args_with_ind:
+            if isinstance(arg_with_ind.element, Identity):
+                k = arg_with_ind.element.shape[0]
+                # Candidate for removal:
+                if arg_with_ind.indices == [None, None]:
+                    # Free identity matrix, will be cleared by _remove_trivial_dims:
+                    continue
+                elif None in arg_with_ind.indices:
+                    ind = [j for j in arg_with_ind.indices if j is not None][0]
+                    counted = editor.count_args_with_index(ind)
+                    if counted == 1:
+                        # Identity matrix contracted only on one index with itself,
+                        # transform to a OneArray(k) element:
+                        editor.insert_after(arg_with_ind, OneArray(k))
+                        editor.args_with_ind.remove(arg_with_ind)
+                        flag = True
+                        break
+                    elif counted > 2:
+                        # Case counted = 2 is a matrix multiplication by identity matrix, skip it.
+                        # Case counted > 2 is a multiple contraction,
+                        # this is a case where the contraction becomes a diagonalization if the
+                        # identity matrix is dropped.
+                        continue
+                elif arg_with_ind.indices[0] == arg_with_ind.indices[1]:
+                    ind = arg_with_ind.indices[0]
+                    counted = editor.count_args_with_index(ind)
+                    if counted > 1:
+                        editor.args_with_ind.remove(arg_with_ind)
+                        flag = True
+                        break
+                    else:
+                        # This is a trace, skip it as it will be recognized somewhere else:
+                        pass
+            elif ask(Q.diagonal(arg_with_ind.element)):
+                if arg_with_ind.indices == [None, None]:
+                    continue
+                elif None in arg_with_ind.indices:
+                    pass
+                elif arg_with_ind.indices[0] == arg_with_ind.indices[1]:
+                    ind = arg_with_ind.indices[0]
+                    counted = editor.count_args_with_index(ind)
+                    if counted == 3:
+                        # A_ai B_bi D_ii ==> A_ai D_ij B_bj
+                        ind_new = editor.get_new_contraction_index()
+                        other_args = [j for j in editor.args_with_ind if j != arg_with_ind]
+                        other_args[1].indices = [ind_new if j == ind else j for j in other_args[1].indices]
+                        arg_with_ind.indices = [ind, ind_new]
+                        flag = True
+                        break
+
+    return editor.to_array_contraction()
+
+
+def remove_identity_matrices(expr: ArrayContraction):
+    editor = _EditArrayContraction(expr)
+    removed: List[int] = []
+
+    permutation_map = {}
+
+    free_indices = list(accumulate([0] + [sum(i is None for i in arg.indices) for arg in editor.args_with_ind]))
+    free_map = dict(zip(editor.args_with_ind, free_indices[:-1]))
+
+    update_pairs = {}
+
+    for ind in range(editor.number_of_contraction_indices):
+        args = editor.get_args_with_index(ind)
+        identity_matrices = [i for i in args if isinstance(i.element, Identity)]
+        number_identity_matrices = len(identity_matrices)
+        # If the contraction involves a non-identity matrix and multiple identity matrices:
+        if number_identity_matrices != len(args) - 1 or number_identity_matrices == 0:
+            continue
+        # Get the non-identity element:
+        non_identity = [i for i in args if not isinstance(i.element, Identity)][0]
+        # Check that all identity matrices have at least one free index
+        # (otherwise they would be contractions to some other elements)
+        if any(None not in i.indices for i in identity_matrices):
+            continue
+        # Mark the identity matrices for removal:
+        for i in identity_matrices:
+            i.element = None
+            removed.extend(range(free_map[i], free_map[i] + len([j for j in i.indices if j is None])))
+        last_removed = removed.pop(-1)
+        update_pairs[last_removed, ind] = non_identity.indices[:]
+        # Remove the indices from the non-identity matrix, as the contraction
+        # no longer exists:
+        non_identity.indices = [None if i == ind else i for i in non_identity.indices]
+
+    removed.sort()
+
+    shifts = list(accumulate([1 if i in removed else 0 for i in range(get_rank(expr))]))
+    for (last_removed, ind), non_identity_indices in update_pairs.items():
+        pos = [free_map[non_identity] + i for i, e in enumerate(non_identity_indices) if e == ind]
+        assert len(pos) == 1
+        for j in pos:
+            permutation_map[j] = last_removed
+
+    editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None]
+    ret_expr = editor.to_array_contraction()
+    permutation = []
+    counter = 0
+    counter2 = 0
+    for j in range(get_rank(expr)):
+        if j in removed:
+            continue
+        if counter2 in permutation_map:
+            target = permutation_map[counter2]
+            permutation.append(target - shifts[target])
+            counter2 += 1
+        else:
+            while counter in permutation_map.values():
+                counter += 1
+            permutation.append(counter)
+            counter += 1
+            counter2 += 1
+    ret_expr2 = _permute_dims(ret_expr, _af_invert(permutation))
+    return ret_expr2, removed
+
+
+def _combine_removed(dim: int, removed1: List[int], removed2: List[int]) -> List[int]:
+    # Concatenate two axis removal operations as performed by
+    # _remove_trivial_dims,
+    removed1 = sorted(removed1)
+    removed2 = sorted(removed2)
+    i = 0
+    j = 0
+    removed = []
+    while True:
+        if j >= len(removed2):
+            while i < len(removed1):
+                removed.append(removed1[i])
+                i += 1
+            break
+        elif i < len(removed1) and removed1[i] <= i + removed2[j]:
+            removed.append(removed1[i])
+            i += 1
+        else:
+            removed.append(i + removed2[j])
+            j += 1
+    return removed
+
+
+def _array_contraction_to_diagonal_multiple_identity(expr: ArrayContraction):
+    editor = _EditArrayContraction(expr)
+    editor.track_permutation_start()
+    removed: List[int] = []
+    diag_index_counter: int = 0
+    for i in range(editor.number_of_contraction_indices):
+        identities = []
+        args = []
+        for j, arg in enumerate(editor.args_with_ind):
+            if i not in arg.indices:
+                continue
+            if isinstance(arg.element, Identity):
+                identities.append(arg)
+            else:
+                args.append(arg)
+        if len(identities) == 0:
+            continue
+        if len(args) + len(identities) < 3:
+            continue
+        new_diag_ind = -1 - diag_index_counter
+        diag_index_counter += 1
+        # Variable "flag" to control whether to skip this contraction set:
+        flag: bool = True
+        for i1, id1 in enumerate(identities):
+            if None not in id1.indices:
+                flag = True
+                break
+            free_pos = list(range(*editor.get_absolute_free_range(id1)))[0]
+            editor._track_permutation[-1].append(free_pos) # type: ignore
+            id1.element = None
+            flag = False
+            break
+        if flag:
+            continue
+        for arg in identities[:i1] + identities[i1+1:]:
+            arg.element = None
+            removed.extend(range(*editor.get_absolute_free_range(arg)))
+        for arg in args:
+            arg.indices = [new_diag_ind if j == i else j for j in arg.indices]
+    for j, e in enumerate(editor.args_with_ind):
+        if e.element is None:
+            editor._track_permutation[j] = None # type: ignore
+    editor._track_permutation = [i for i in editor._track_permutation if i is not None] # type: ignore
+    # Renumber permutation array form in order to deal with deleted positions:
+    remap = {e: i for i, e in enumerate(sorted({k for j in editor._track_permutation for k in j}))}
+    editor._track_permutation = [[remap[j] for j in i] for i in editor._track_permutation]
+    editor.args_with_ind = [i for i in editor.args_with_ind if i.element is not None]
+    new_expr = editor.to_array_contraction()
+    return new_expr, removed

.venv/lib/python3.11/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, ()

.venv/lib/python3.11/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

.venv/lib/python3.11/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)