add read me

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_arrayexpr_derivatives.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_as_explicit.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_indexed.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_convert_array_to_matrix.py

@@ -0,0 +1,689 @@
+from sympy import Lambda, S, Dummy, KroneckerProduct
+from sympy.core.symbol import symbols
+from sympy.functions.elementary.miscellaneous import sqrt
+from sympy.functions.elementary.trigonometric import cos, sin
+from sympy.matrices.expressions.hadamard import HadamardProduct, HadamardPower
+from sympy.matrices.expressions.special import (Identity, OneMatrix, ZeroMatrix)
+from sympy.matrices.expressions.matexpr import MatrixElement
+from sympy.tensor.array.expressions.from_matrix_to_array import convert_matrix_to_array
+from sympy.tensor.array.expressions.from_array_to_matrix import _support_function_tp1_recognize, \
+    _array_diag2contr_diagmatrix, convert_array_to_matrix, _remove_trivial_dims, _array2matrix, \
+    _combine_removed, identify_removable_identity_matrices, _array_contraction_to_diagonal_multiple_identity
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.combinatorics import Permutation
+from sympy.matrices.expressions.diagonal import DiagMatrix, DiagonalMatrix
+from sympy.matrices import Trace, MatMul, Transpose
+from sympy.tensor.array.expressions.array_expressions import ZeroArray, OneArray, \
+    ArrayElement, ArraySymbol, ArrayElementwiseApplyFunc, _array_tensor_product, _array_contraction, \
+    _array_diagonal, _permute_dims, PermuteDims, ArrayAdd, ArrayDiagonal, ArrayContraction, ArrayTensorProduct
+from sympy.testing.pytest import raises
+
+
+i, j, k, l, m, n = symbols("i j k l m n")
+
+I = Identity(k)
+I1 = Identity(1)
+
+M = MatrixSymbol("M", k, k)
+N = MatrixSymbol("N", k, k)
+P = MatrixSymbol("P", k, k)
+Q = MatrixSymbol("Q", k, k)
+
+A = MatrixSymbol("A", k, k)
+B = MatrixSymbol("B", k, k)
+C = MatrixSymbol("C", k, k)
+D = MatrixSymbol("D", k, k)
+
+X = MatrixSymbol("X", k, k)
+Y = MatrixSymbol("Y", k, k)
+
+a = MatrixSymbol("a", k, 1)
+b = MatrixSymbol("b", k, 1)
+c = MatrixSymbol("c", k, 1)
+d = MatrixSymbol("d", k, 1)
+
+x = MatrixSymbol("x", k, 1)
+y = MatrixSymbol("y", k, 1)
+
+
+def test_arrayexpr_convert_array_to_matrix():
+
+    cg = _array_contraction(_array_tensor_product(M), (0, 1))
+    assert convert_array_to_matrix(cg) == Trace(M)
+
+    cg = _array_contraction(_array_tensor_product(M, N), (0, 1), (2, 3))
+    assert convert_array_to_matrix(cg) == Trace(M) * Trace(N)
+
+    cg = _array_contraction(_array_tensor_product(M, N), (0, 3), (1, 2))
+    assert convert_array_to_matrix(cg) == Trace(M * N)
+
+    cg = _array_contraction(_array_tensor_product(M, N), (0, 2), (1, 3))
+    assert convert_array_to_matrix(cg) == Trace(M * N.T)
+
+    cg = convert_matrix_to_array(M * N * P)
+    assert convert_array_to_matrix(cg) == M * N * P
+
+    cg = convert_matrix_to_array(M * N.T * P)
+    assert convert_array_to_matrix(cg) == M * N.T * P
+
+    cg = _array_contraction(_array_tensor_product(M,N,P,Q), (1, 2), (5, 6))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M * N, P * Q)
+
+    cg = _array_contraction(_array_tensor_product(-2, M, N), (1, 2))
+    assert convert_array_to_matrix(cg) == -2 * M * N
+
+    a = MatrixSymbol("a", k, 1)
+    b = MatrixSymbol("b", k, 1)
+    c = MatrixSymbol("c", k, 1)
+    cg = PermuteDims(
+        _array_contraction(
+            _array_tensor_product(
+                a,
+                ArrayAdd(
+                    _array_tensor_product(b, c),
+                    _array_tensor_product(c, b),
+                )
+            ), (2, 4)), [0, 1, 3, 2])
+    assert convert_array_to_matrix(cg) == a * (b.T * c + c.T * b)
+
+    za = ZeroArray(m, n)
+    assert convert_array_to_matrix(za) == ZeroMatrix(m, n)
+
+    cg = _array_tensor_product(3, M)
+    assert convert_array_to_matrix(cg) == 3 * M
+
+    # Partial conversion to matrix multiplication:
+    expr = _array_contraction(_array_tensor_product(M, N, P, Q), (0, 2), (1, 4, 6))
+    assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(M.T*N, P, Q), (0, 2, 4))
+
+    x = MatrixSymbol("x", k, 1)
+    cg = PermuteDims(
+        _array_contraction(_array_tensor_product(OneArray(1), x, OneArray(1), DiagMatrix(Identity(1))),
+                                (0, 5)), Permutation(1, 2, 3))
+    assert convert_array_to_matrix(cg) == x
+
+    expr = ArrayAdd(M, PermuteDims(M, [1, 0]))
+    assert convert_array_to_matrix(expr) == M + Transpose(M)
+
+
+def test_arrayexpr_convert_array_to_matrix2():
+    cg = _array_contraction(_array_tensor_product(M, N), (1, 3))
+    assert convert_array_to_matrix(cg) == M * N.T
+
+    cg = PermuteDims(_array_tensor_product(M, N), Permutation([0, 1, 3, 2]))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
+
+    cg = _array_tensor_product(M, PermuteDims(N, Permutation([1, 0])))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
+
+    cg = _array_contraction(
+        PermuteDims(
+            _array_tensor_product(M, N, P, Q), Permutation([0, 2, 3, 1, 4, 5, 7, 6])),
+        (1, 2), (3, 5)
+    )
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
+
+    cg = _array_contraction(
+        _array_tensor_product(M, N, P, PermuteDims(Q, Permutation([1, 0]))),
+        (1, 5), (2, 3)
+    )
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M * P.T * Trace(N), Q.T)
+
+    cg = _array_tensor_product(M, PermuteDims(N, [1, 0]))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M, N.T)
+
+    cg = _array_tensor_product(PermuteDims(M, [1, 0]), PermuteDims(N, [1, 0]))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(M.T, N.T)
+
+    cg = _array_tensor_product(PermuteDims(N, [1, 0]), PermuteDims(M, [1, 0]))
+    assert convert_array_to_matrix(cg) == _array_tensor_product(N.T, M.T)
+
+    cg = _array_contraction(M, (0,), (1,))
+    assert convert_array_to_matrix(cg) == OneMatrix(1, k)*M*OneMatrix(k, 1)
+
+    cg = _array_contraction(x, (0,), (1,))
+    assert convert_array_to_matrix(cg) == OneMatrix(1, k)*x
+
+    Xm = MatrixSymbol("Xm", m, n)
+    cg = _array_contraction(Xm, (0,), (1,))
+    assert convert_array_to_matrix(cg) == OneMatrix(1, m)*Xm*OneMatrix(n, 1)
+
+
+def test_arrayexpr_convert_array_to_diagonalized_vector():
+
+    # Check matrix recognition over trivial dimensions:
+
+    cg = _array_tensor_product(a, b)
+    assert convert_array_to_matrix(cg) == a * b.T
+
+    cg = _array_tensor_product(I1, a, b)
+    assert convert_array_to_matrix(cg) == a * b.T
+
+    # Recognize trace inside a tensor product:
+
+    cg = _array_contraction(_array_tensor_product(A, B, C), (0, 3), (1, 2))
+    assert convert_array_to_matrix(cg) == Trace(A * B) * C
+
+    # Transform diagonal operator to contraction:
+
+    cg = _array_diagonal(_array_tensor_product(A, a), (1, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (1, 3))
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(a, b), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _permute_dims(
+        _array_contraction(_array_tensor_product(DiagMatrix(a), OneArray(1), b), (0, 3)), [1, 2, 0]
+    )
+    assert convert_array_to_matrix(cg) == b.T * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(A, a), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(A, OneArray(1), DiagMatrix(a)), (0, 3))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a)
+
+    cg = _array_diagonal(_array_tensor_product(I, x, I1), (0, 2), (3, 5))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(I, OneArray(1), I1, DiagMatrix(x)), (0, 5))
+    assert convert_array_to_matrix(cg) == DiagMatrix(x)
+
+    cg = _array_diagonal(_array_tensor_product(I, x, A, B), (1, 2), (5, 6))
+    assert _array_diag2contr_diagmatrix(cg) == _array_diagonal(_array_contraction(_array_tensor_product(I, OneArray(1), A, B, DiagMatrix(x)), (1, 7)), (5, 6))
+    # TODO: this is returning a wrong result:
+    # convert_array_to_matrix(cg)
+
+    cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3, 5))
+    assert convert_array_to_matrix(cg) == a*b.T
+
+    cg = _array_diagonal(_array_tensor_product(I1, a, b), (1, 3))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), a, b, I1), (2, 6))
+    assert convert_array_to_matrix(cg) == a*b.T
+
+    cg = _array_diagonal(_array_tensor_product(x, I1), (1, 2))
+    assert isinstance(cg, ArrayDiagonal)
+    assert cg.diagonal_indices == ((1, 2),)
+    assert convert_array_to_matrix(cg) == x
+
+    cg = _array_diagonal(_array_tensor_product(x, I), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg) == _array_contraction(_array_tensor_product(OneArray(1), I, DiagMatrix(x)), (1, 3))
+    assert convert_array_to_matrix(cg).doit() == DiagMatrix(x)
+
+    raises(ValueError, lambda: _array_diagonal(x, (1,)))
+
+    # Ignore identity matrices with contractions:
+
+    cg = _array_contraction(_array_tensor_product(I, A, I, I), (0, 2), (1, 3), (5, 7))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg) == Trace(A) * I
+
+    cg = _array_contraction(_array_tensor_product(Trace(A) * I, I, I), (1, 5), (3, 4))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg).doit() == Trace(A) * I
+
+    # Add DiagMatrix when required:
+
+    cg = _array_contraction(_array_tensor_product(A, a), (1, 2))
+    assert cg.split_multiple_contractions() == cg
+    assert convert_array_to_matrix(cg) == A * a
+
+    cg = _array_contraction(_array_tensor_product(A, a, B), (1, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (1, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(A, a, B), (0, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), B), (0, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(A, a, b, a.T, B), (0, 2, 4, 7, 9))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1),
+                                                DiagMatrix(b), OneArray(1), DiagMatrix(a), OneArray(1), B),
+                                               (0, 2), (3, 5), (6, 9), (8, 12))
+    assert convert_array_to_matrix(cg) == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T
+
+    cg = _array_contraction(_array_tensor_product(I1, I1, I1), (1, 2, 4))
+    assert cg.split_multiple_contractions() == _array_contraction(_array_tensor_product(I1, I1, OneArray(1), I1), (1, 2), (3, 5))
+    assert convert_array_to_matrix(cg) == 1
+
+    cg = _array_contraction(_array_tensor_product(I, I, I, I, A), (1, 2, 8), (5, 6, 9))
+    assert convert_array_to_matrix(cg.split_multiple_contractions()).doit() == A
+
+    cg = _array_contraction(_array_tensor_product(A, a, C, a, B), (1, 2, 4), (5, 6, 8))
+    expected = _array_contraction(_array_tensor_product(A, DiagMatrix(a), OneArray(1), C, DiagMatrix(a), OneArray(1), B), (1, 3), (2, 5), (6, 7), (8, 10))
+    assert cg.split_multiple_contractions() == expected
+    assert convert_array_to_matrix(cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B
+
+    cg = _array_contraction(_array_tensor_product(a, I1, b, I1, (a.T*b).applyfunc(cos)), (1, 2, 8), (5, 6, 9))
+    expected = _array_contraction(_array_tensor_product(a, I1, OneArray(1), b, I1, OneArray(1), (a.T*b).applyfunc(cos)),
+                                (1, 3), (2, 10), (6, 8), (7, 11))
+    assert cg.split_multiple_contractions().dummy_eq(expected)
+    assert convert_array_to_matrix(cg).doit().dummy_eq(MatMul(a, (a.T * b).applyfunc(cos), b.T))
+
+
+def test_arrayexpr_convert_array_contraction_tp_additions():
+    a = ArrayAdd(
+        _array_tensor_product(M, N),
+        _array_tensor_product(N, M)
+    )
+    tp = _array_tensor_product(P, a, Q)
+    expr = _array_contraction(tp, (3, 4))
+    expected = _array_tensor_product(
+        P,
+        ArrayAdd(
+            _array_contraction(_array_tensor_product(M, N), (1, 2)),
+            _array_contraction(_array_tensor_product(N, M), (1, 2)),
+        ),
+        Q
+    )
+    assert expr == expected
+    assert convert_array_to_matrix(expr) == _array_tensor_product(P, M * N + N * M, Q)
+
+    expr = _array_contraction(tp, (1, 2), (3, 4), (5, 6))
+    result = _array_contraction(
+        _array_tensor_product(
+            P,
+            ArrayAdd(
+                _array_contraction(_array_tensor_product(M, N), (1, 2)),
+                _array_contraction(_array_tensor_product(N, M), (1, 2)),
+            ),
+            Q
+        ), (1, 2), (3, 4))
+    assert expr == result
+    assert convert_array_to_matrix(expr) == P * (M * N + N * M) * Q
+
+
+def test_arrayexpr_convert_array_to_implicit_matmul():
+    # Trivial dimensions are suppressed, so the result can be expressed in matrix form:
+
+    cg = _array_tensor_product(a, b)
+    assert convert_array_to_matrix(cg) == a * b.T
+
+    cg = _array_tensor_product(a, b, I)
+    assert convert_array_to_matrix(cg) == _array_tensor_product(a*b.T, I)
+
+    cg = _array_tensor_product(I, a, b)
+    assert convert_array_to_matrix(cg) == _array_tensor_product(I, a*b.T)
+
+    cg = _array_tensor_product(a, I, b)
+    assert convert_array_to_matrix(cg) == _array_tensor_product(a, I, b)
+
+    cg = _array_contraction(_array_tensor_product(I, I), (1, 2))
+    assert convert_array_to_matrix(cg) == I
+
+    cg = PermuteDims(_array_tensor_product(I, Identity(1)), [0, 2, 1, 3])
+    assert convert_array_to_matrix(cg) == I
+
+
+def test_arrayexpr_convert_array_to_matrix_remove_trivial_dims():
+
+    # Tensor Product:
+    assert _remove_trivial_dims(_array_tensor_product(a, b)) == (a * b.T, [1, 3])
+    assert _remove_trivial_dims(_array_tensor_product(a.T, b)) == (a * b.T, [0, 3])
+    assert _remove_trivial_dims(_array_tensor_product(a, b.T)) == (a * b.T, [1, 2])
+    assert _remove_trivial_dims(_array_tensor_product(a.T, b.T)) == (a * b.T, [0, 2])
+
+    assert _remove_trivial_dims(_array_tensor_product(I, a.T, b.T)) == (_array_tensor_product(I, a * b.T), [2, 4])
+    assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T)) == (_array_tensor_product(a.T, I, b.T), [])
+
+    assert _remove_trivial_dims(_array_tensor_product(a, I)) == (_array_tensor_product(a, I), [])
+    assert _remove_trivial_dims(_array_tensor_product(I, a)) == (_array_tensor_product(I, a), [])
+
+    assert _remove_trivial_dims(_array_tensor_product(a.T, b.T, c, d)) == (
+        _array_tensor_product(a * b.T, c * d.T), [0, 2, 5, 7])
+    assert _remove_trivial_dims(_array_tensor_product(a.T, I, b.T, c, d, I)) == (
+        _array_tensor_product(a.T, I, b*c.T, d, I), [4, 7])
+
+    # Addition:
+
+    cg = ArrayAdd(_array_tensor_product(a, b), _array_tensor_product(c, d))
+    assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
+
+    # Permute Dims:
+
+    cg = PermuteDims(_array_tensor_product(a, b), Permutation(3)(1, 2))
+    assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
+
+    cg = PermuteDims(_array_tensor_product(a, I, b), Permutation(5)(1, 2, 3, 4))
+    assert _remove_trivial_dims(cg) == (cg, [])
+
+    cg = PermuteDims(_array_tensor_product(I, b, a), Permutation(5)(1, 2, 4, 5, 3))
+    assert _remove_trivial_dims(cg) == (PermuteDims(_array_tensor_product(I, b * a.T), [0, 2, 3, 1]), [4, 5])
+
+    # Diagonal:
+
+    cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
+    assert _remove_trivial_dims(cg) == (cg, [])
+
+    # Contraction:
+
+    cg = _array_contraction(_array_tensor_product(M, a), (1, 2))
+    assert _remove_trivial_dims(cg) == (cg, [])
+
+    # A few more cases to test the removal and shift of nested removed axes
+    # with array contractions and array diagonals:
+    tp = _array_tensor_product(
+        OneMatrix(1, 1),
+        M,
+        x,
+        OneMatrix(1, 1),
+        Identity(1),
+    )
+
+    expr = _array_contraction(tp, (1, 8))
+    rexpr, removed = _remove_trivial_dims(expr)
+    assert removed == [0, 5, 6, 7]
+
+    expr = _array_contraction(tp, (1, 8), (3, 4))
+    rexpr, removed = _remove_trivial_dims(expr)
+    assert removed == [0, 3, 4, 5]
+
+    expr = _array_diagonal(tp, (1, 8))
+    rexpr, removed = _remove_trivial_dims(expr)
+    assert removed == [0, 5, 6, 7, 8]
+
+    expr = _array_diagonal(tp, (1, 8), (3, 4))
+    rexpr, removed = _remove_trivial_dims(expr)
+    assert removed == [0, 3, 4, 5, 6]
+
+    expr = _array_diagonal(_array_contraction(_array_tensor_product(A, x, I, I1), (1, 2, 5)), (1, 4))
+    rexpr, removed = _remove_trivial_dims(expr)
+    assert removed == [2, 3]
+
+    cg = _array_diagonal(_array_tensor_product(PermuteDims(_array_tensor_product(x, I1), Permutation(1, 2, 3)), (x.T*x).applyfunc(sqrt)), (2, 4), (3, 5))
+    rexpr, removed = _remove_trivial_dims(cg)
+    assert removed == [1, 2]
+
+    # Contractions with identity matrices need to be followed by a permutation
+    # in order
+    cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8))
+    ret, removed = _remove_trivial_dims(cg)
+    assert ret == PermuteDims(_array_tensor_product(A, B, C, M), [0, 2, 3, 4, 5, 6, 7, 1])
+    assert removed == []
+
+    cg = _array_contraction(_array_tensor_product(A, B, C, M, I), (1, 8), (3, 4))
+    ret, removed = _remove_trivial_dims(cg)
+    assert ret == PermuteDims(_array_contraction(_array_tensor_product(A, B, C, M), (3, 4)), [0, 2, 3, 4, 5, 1])
+    assert removed == []
+
+    # Trivial matrices are sometimes inserted into MatMul expressions:
+
+    cg = _array_tensor_product(b*b.T, a.T*a)
+    ret, removed = _remove_trivial_dims(cg)
+    assert ret == b*a.T*a*b.T
+    assert removed == [2, 3]
+
+    Xs = ArraySymbol("X", (3, 2, k))
+    cg = _array_tensor_product(M, Xs, b.T*c, a*a.T, b*b.T, c.T*d)
+    ret, removed = _remove_trivial_dims(cg)
+    assert ret == _array_tensor_product(M, Xs, a*b.T*c*c.T*d*a.T, b*b.T)
+    assert removed == [5, 6, 11, 12]
+
+    cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
+    assert _remove_trivial_dims(cg) == (PermuteDims(_array_diagonal(_array_tensor_product(I, x), (1, 2)), Permutation(1, 2)), [1])
+
+    expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2))
+    assert _remove_trivial_dims(expr) == (PermuteDims(_array_tensor_product(DiagMatrix(x), y), [1, 2, 3, 0]), [0])
+
+    expr = _array_diagonal(_array_tensor_product(x, I, y), (0, 2), (3, 4))
+    assert _remove_trivial_dims(expr) == (expr, [])
+
+
+def test_arrayexpr_convert_array_to_matrix_diag2contraction_diagmatrix():
+    cg = _array_diagonal(_array_tensor_product(M, a), (1, 2))
+    res = _array_diag2contr_diagmatrix(cg)
+    assert res.shape == cg.shape
+    assert res == _array_contraction(_array_tensor_product(M, OneArray(1), DiagMatrix(a)), (1, 3))
+
+    raises(ValueError, lambda: _array_diagonal(_array_tensor_product(a, M), (1, 2)))
+
+    cg = _array_diagonal(_array_tensor_product(a.T, M), (1, 2))
+    res = _array_diag2contr_diagmatrix(cg)
+    assert res.shape == cg.shape
+    assert res == _array_contraction(_array_tensor_product(OneArray(1), M, DiagMatrix(a.T)), (1, 4))
+
+    cg = _array_diagonal(_array_tensor_product(a.T, M, N, b.T), (1, 2), (4, 7))
+    res = _array_diag2contr_diagmatrix(cg)
+    assert res.shape == cg.shape
+    assert res == _array_contraction(
+        _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9))
+
+    cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 2), (4, 7))
+    res = _array_diag2contr_diagmatrix(cg)
+    assert res.shape == cg.shape
+    assert res == _array_contraction(
+        _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9))
+
+    cg = _array_diagonal(_array_tensor_product(a, M, N, b.T), (0, 4), (3, 7))
+    res = _array_diag2contr_diagmatrix(cg)
+    assert res.shape == cg.shape
+    assert res == _array_contraction(
+        _array_tensor_product(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9))
+
+    I1 = Identity(1)
+    x = MatrixSymbol("x", k, 1)
+    A = MatrixSymbol("A", k, k)
+    cg = _array_diagonal(_array_tensor_product(x, A.T, I1), (0, 2))
+    assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
+    assert _array2matrix(cg).shape == cg.shape
+
+
+def test_arrayexpr_convert_array_to_matrix_support_function():
+
+    assert _support_function_tp1_recognize([], [2 * k]) == 2 * k
+
+    assert _support_function_tp1_recognize([(1, 2)], [A, 2 * k, B, 3]) == 6 * k * A * B
+
+    assert _support_function_tp1_recognize([(0, 3), (1, 2)], [A, B]) == Trace(A * B)
+
+    assert _support_function_tp1_recognize([(1, 2)], [A, B]) == A * B
+    assert _support_function_tp1_recognize([(0, 2)], [A, B]) == A.T * B
+    assert _support_function_tp1_recognize([(1, 3)], [A, B]) == A * B.T
+    assert _support_function_tp1_recognize([(0, 3)], [A, B]) == A.T * B.T
+
+    assert _support_function_tp1_recognize([(1, 2), (5, 6)], [A, B, C, D]) == _array_tensor_product(A * B, C * D)
+    assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
+        _array_tensor_product(A * C, B * D), [0, 2, 1, 3])
+
+    assert _support_function_tp1_recognize([(0, 3), (1, 4)], [A, B, C]) == B * A * C
+
+    assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4), (7, 8)],
+                                           [X, Y, A, B, C, D]) == X * Y * A * B * C * D
+
+    assert _support_function_tp1_recognize([(9, 10), (1, 2), (5, 6), (3, 4)],
+                                           [X, Y, A, B, C, D]) == _array_tensor_product(X * Y * A * B, C * D)
+
+    assert _support_function_tp1_recognize([(1, 7), (3, 8), (4, 11)], [X, Y, A, B, C, D]) == PermuteDims(
+        _array_tensor_product(X * B.T, Y * C, A.T * D.T), [0, 2, 4, 1, 3, 5]
+    )
+
+    assert _support_function_tp1_recognize([(0, 1), (3, 6), (5, 8)], [X, A, B, C, D]) == PermuteDims(
+        _array_tensor_product(Trace(X) * A * C, B * D), [0, 2, 1, 3])
+
+    assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [A, A, B, C, D]) == A ** 2 * B * C * D
+    assert _support_function_tp1_recognize([(1, 2), (3, 4), (5, 6), (7, 8)], [X, A, B, C, D]) == X * A * B * C * D
+
+    assert _support_function_tp1_recognize([(1, 6), (3, 8), (5, 10)], [X, Y, A, B, C, D]) == PermuteDims(
+        _array_tensor_product(X * B, Y * C, A * D), [0, 2, 4, 1, 3, 5]
+    )
+
+    assert _support_function_tp1_recognize([(1, 4), (3, 6)], [A, B, C, D]) == PermuteDims(
+        _array_tensor_product(A * C, B * D), [0, 2, 1, 3])
+
+    assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
+
+    assert _support_function_tp1_recognize([(0, 4), (1, 7), (2, 5), (3, 8)], [X, A, B, C, D]) == C*X.T*B*A*D
+
+
+def test_convert_array_to_hadamard_products():
+
+    expr = HadamardProduct(M, N)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = HadamardProduct(M, N)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = Q*HadamardProduct(M, N)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = Q*HadamardProduct(M, N.T)*P
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == expr
+
+    expr = HadamardProduct(M, N)*HadamardProduct(Q, P)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert expr == ret
+
+    expr = P.T*HadamardProduct(M, N)*HadamardProduct(Q, P)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert expr == ret
+
+    # ArrayDiagonal should be converted
+    cg = _array_diagonal(_array_tensor_product(M, N, Q), (1, 3), (0, 2, 4))
+    ret = convert_array_to_matrix(cg)
+    expected = PermuteDims(_array_diagonal(_array_tensor_product(HadamardProduct(M.T, N.T), Q), (1, 2)), [1, 0, 2])
+    assert expected == ret
+
+    # Special case that should return the same expression:
+    cg = _array_diagonal(_array_tensor_product(HadamardProduct(M, N), Q), (0, 2))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    # Hadamard products with traces:
+
+    expr = Trace(HadamardProduct(M, N))
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M.T, N.T))
+
+    expr = Trace(A*HadamardProduct(M, N))
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M, N)*A)
+
+    expr = Trace(HadamardProduct(A, M)*N)
+    cg = convert_matrix_to_array(expr)
+    ret = convert_array_to_matrix(cg)
+    assert ret == Trace(HadamardProduct(M.T, N)*A)
+
+    # These should not be converted into Hadamard products:
+
+    cg = _array_diagonal(_array_tensor_product(M, N), (0, 1, 2, 3))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    cg = _array_diagonal(_array_tensor_product(A), (0, 1))
+    ret = convert_array_to_matrix(cg)
+    assert ret == cg
+
+    cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 2, 4), (1, 3, 5))
+    assert convert_array_to_matrix(cg) == HadamardProduct(M, N, P)
+
+    cg = _array_diagonal(_array_tensor_product(M, N, P), (0, 3, 4), (1, 2, 5))
+    assert convert_array_to_matrix(cg) == HadamardProduct(M, P, N.T)
+
+    cg = _array_diagonal(_array_tensor_product(I, I1, x), (1, 4), (3, 5))
+    assert convert_array_to_matrix(cg) == DiagMatrix(x)
+
+
+def test_identify_removable_identity_matrices():
+
+    D = DiagonalMatrix(MatrixSymbol("D", k, k))
+
+    cg = _array_contraction(_array_tensor_product(A, B, I), (1, 2, 4, 5))
+    expected = _array_contraction(_array_tensor_product(A, B), (1, 2))
+    assert identify_removable_identity_matrices(cg) == expected
+
+    cg = _array_contraction(_array_tensor_product(A, B, C, I), (1, 3, 5, 6, 7))
+    expected = _array_contraction(_array_tensor_product(A, B, C), (1, 3, 5))
+    assert identify_removable_identity_matrices(cg) == expected
+
+    # Tests with diagonal matrices:
+
+    cg = _array_contraction(_array_tensor_product(A, B, D), (1, 2, 4, 5))
+    ret = identify_removable_identity_matrices(cg)
+    expected = _array_contraction(_array_tensor_product(A, B, D), (1, 4), (2, 5))
+    assert ret == expected
+
+    cg = _array_contraction(_array_tensor_product(A, B, D, M, N), (1, 2, 4, 5, 6, 8))
+    ret = identify_removable_identity_matrices(cg)
+    assert ret == cg
+
+
+def test_combine_removed():
+
+    assert _combine_removed(6, [0, 1, 2], [0, 1, 2]) == [0, 1, 2, 3, 4, 5]
+    assert _combine_removed(8, [2, 5], [1, 3, 4]) == [1, 2, 4, 5, 6]
+    assert _combine_removed(8, [7], []) == [7]
+
+
+def test_array_contraction_to_diagonal_multiple_identities():
+
+    expr = _array_contraction(_array_tensor_product(A, B, I, C), (1, 2, 4), (5, 6))
+    assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
+    assert convert_array_to_matrix(expr) == _array_contraction(_array_tensor_product(A, B, C), (1, 2, 4))
+
+    expr = _array_contraction(_array_tensor_product(A, I, I), (1, 2, 4))
+    assert _array_contraction_to_diagonal_multiple_identity(expr) == (A, [2])
+    assert convert_array_to_matrix(expr) == A
+
+    expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 4), (3, 6))
+    assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
+
+    expr = _array_contraction(_array_tensor_product(A, I, I, B), (1, 2, 3, 4, 6))
+    assert _array_contraction_to_diagonal_multiple_identity(expr) == (expr, [])
+
+
+def test_convert_array_element_to_matrix():
+
+    expr = ArrayElement(M, (i, j))
+    assert convert_array_to_matrix(expr) == MatrixElement(M, i, j)
+
+    expr = ArrayElement(_array_contraction(_array_tensor_product(M, N), (1, 3)), (i, j))
+    assert convert_array_to_matrix(expr) == MatrixElement(M*N.T, i, j)
+
+    expr = ArrayElement(_array_tensor_product(M, N), (i, j, m, n))
+    assert convert_array_to_matrix(expr) == expr
+
+
+def test_convert_array_elementwise_function_to_matrix():
+
+    d = Dummy("d")
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x.T*y)
+    assert convert_array_to_matrix(expr) == sin(x.T*y)
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, d**2), x.T*y)
+    assert convert_array_to_matrix(expr) == (x.T*y)**2
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, sin(d)), x)
+    assert convert_array_to_matrix(expr).dummy_eq(x.applyfunc(sin))
+
+    expr = ArrayElementwiseApplyFunc(Lambda(d, 1 / (2 * sqrt(d))), x)
+    assert convert_array_to_matrix(expr) == S.Half * HadamardPower(x, -S.Half)
+
+
+def test_array2matrix():
+    # See issue https://github.com/sympy/sympy/pull/22877
+    expr = PermuteDims(ArrayContraction(ArrayTensorProduct(x, I, I1, x), (0, 3), (1, 7)), Permutation(2, 3))
+    expected = PermuteDims(ArrayTensorProduct(x*x.T, I1), Permutation(3)(1, 2))
+    assert _array2matrix(expr) == expected
+
+
+def test_recognize_broadcasting():
+    expr = ArrayTensorProduct(x.T*x, A)
+    assert _remove_trivial_dims(expr) == (KroneckerProduct(x.T*x, A), [0, 1])
+
+    expr = ArrayTensorProduct(A, x.T*x)
+    assert _remove_trivial_dims(expr) == (KroneckerProduct(A, x.T*x), [2, 3])
+
+    expr = ArrayTensorProduct(A, B, x.T*x, C)
+    assert _remove_trivial_dims(expr) == (ArrayTensorProduct(A, KroneckerProduct(B, x.T*x), C), [4, 5])
+
+    # Always prefer matrix multiplication to Kronecker product, if possible:
+    expr = ArrayTensorProduct(a, b, x.T*x)
+    assert _remove_trivial_dims(expr) == (a*x.T*x*b.T, [1, 3, 4, 5])

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_convert_indexed_to_array.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_convert_matrix_to_array.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/tests/test_deprecated_conv_modules.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/expressions/utils.py

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

venv/lib/python3.12/site-packages/sympy/tensor/array/mutable_ndim_array.py

@@ -0,0 +1,13 @@
+from sympy.tensor.array.ndim_array import NDimArray
+
+
+class MutableNDimArray(NDimArray):
+
+    def as_immutable(self):
+        raise NotImplementedError("abstract method")
+
+    def as_mutable(self):
+        return self
+
+    def _sympy_(self):
+        return self.as_immutable()

venv/lib/python3.12/site-packages/sympy/tensor/array/ndim_array.py

@@ -0,0 +1,601 @@
+from sympy.core.basic import Basic
+from sympy.core.containers import (Dict, Tuple)
+from sympy.core.expr import Expr
+from sympy.core.kind import Kind, NumberKind, UndefinedKind
+from sympy.core.numbers import Integer
+from sympy.core.singleton import S
+from sympy.core.sympify import sympify
+from sympy.external.gmpy import SYMPY_INTS
+from sympy.printing.defaults import Printable
+
+import itertools
+from collections.abc import Iterable
+
+
+class ArrayKind(Kind):
+    """
+    Kind for N-dimensional array in SymPy.
+
+    This kind represents the multidimensional array that algebraic
+    operations are defined. Basic class for this kind is ``NDimArray``,
+    but any expression representing the array can have this.
+
+    Parameters
+    ==========
+
+    element_kind : Kind
+        Kind of the element. Default is :obj:NumberKind `<sympy.core.kind.NumberKind>`,
+        which means that the array contains only numbers.
+
+    Examples
+    ========
+
+    Any instance of array class has ``ArrayKind``.
+
+    >>> from sympy import NDimArray
+    >>> NDimArray([1,2,3]).kind
+    ArrayKind(NumberKind)
+
+    Although expressions representing an array may be not instance of
+    array class, it will have ``ArrayKind`` as well.
+
+    >>> from sympy import Integral
+    >>> from sympy.tensor.array import NDimArray
+    >>> from sympy.abc import x
+    >>> intA = Integral(NDimArray([1,2,3]), x)
+    >>> isinstance(intA, NDimArray)
+    False
+    >>> intA.kind
+    ArrayKind(NumberKind)
+
+    Use ``isinstance()`` to check for ``ArrayKind` without specifying
+    the element kind. Use ``is`` with specifying the element kind.
+
+    >>> from sympy.tensor.array import ArrayKind
+    >>> from sympy.core import NumberKind
+    >>> boolA = NDimArray([True, False])
+    >>> isinstance(boolA.kind, ArrayKind)
+    True
+    >>> boolA.kind is ArrayKind(NumberKind)
+    False
+
+    See Also
+    ========
+
+    shape : Function to return the shape of objects with ``MatrixKind``.
+
+    """
+    def __new__(cls, element_kind=NumberKind):
+        obj = super().__new__(cls, element_kind)
+        obj.element_kind = element_kind
+        return obj
+
+    def __repr__(self):
+        return "ArrayKind(%s)" % self.element_kind
+
+    @classmethod
+    def _union(cls, kinds) -> 'ArrayKind':
+        elem_kinds = {e.kind for e in kinds}
+        if len(elem_kinds) == 1:
+            elemkind, = elem_kinds
+        else:
+            elemkind = UndefinedKind
+        return ArrayKind(elemkind)
+
+
+class NDimArray(Printable):
+    """N-dimensional array.
+
+    Examples
+    ========
+
+    Create an N-dim array of zeros:
+
+    >>> from sympy import MutableDenseNDimArray
+    >>> a = MutableDenseNDimArray.zeros(2, 3, 4)
+    >>> a
+    [[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]]
+
+    Create an N-dim array from a list;
+
+    >>> a = MutableDenseNDimArray([[2, 3], [4, 5]])
+    >>> a
+    [[2, 3], [4, 5]]
+
+    >>> b = MutableDenseNDimArray([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]])
+    >>> b
+    [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]
+
+    Create an N-dim array from a flat list with dimension shape:
+
+    >>> a = MutableDenseNDimArray([1, 2, 3, 4, 5, 6], (2, 3))
+    >>> a
+    [[1, 2, 3], [4, 5, 6]]
+
+    Create an N-dim array from a matrix:
+
+    >>> from sympy import Matrix
+    >>> a = Matrix([[1,2],[3,4]])
+    >>> a
+    Matrix([
+    [1, 2],
+    [3, 4]])
+    >>> b = MutableDenseNDimArray(a)
+    >>> b
+    [[1, 2], [3, 4]]
+
+    Arithmetic operations on N-dim arrays
+
+    >>> a = MutableDenseNDimArray([1, 1, 1, 1], (2, 2))
+    >>> b = MutableDenseNDimArray([4, 4, 4, 4], (2, 2))
+    >>> c = a + b
+    >>> c
+    [[5, 5], [5, 5]]
+    >>> a - b
+    [[-3, -3], [-3, -3]]
+
+    """
+
+    _diff_wrt = True
+    is_scalar = False
+
+    def __new__(cls, iterable, shape=None, **kwargs):
+        from sympy.tensor.array import ImmutableDenseNDimArray
+        return ImmutableDenseNDimArray(iterable, shape, **kwargs)
+
+    def __getitem__(self, index):
+        raise NotImplementedError("A subclass of NDimArray should implement __getitem__")
+
+    def _parse_index(self, index):
+        if isinstance(index, (SYMPY_INTS, Integer)):
+            if index >= self._loop_size:
+                raise ValueError("Only a tuple index is accepted")
+            return index
+
+        if self._loop_size == 0:
+            raise ValueError("Index not valid with an empty array")
+
+        if len(index) != self._rank:
+            raise ValueError('Wrong number of array axes')
+
+        real_index = 0
+        # check if input index can exist in current indexing
+        for i in range(self._rank):
+            if (index[i] >= self.shape[i]) or (index[i] < -self.shape[i]):
+                raise ValueError('Index ' + str(index) + ' out of border')
+            if index[i] < 0:
+                real_index += 1
+            real_index = real_index*self.shape[i] + index[i]
+
+        return real_index
+
+    def _get_tuple_index(self, integer_index):
+        index = []
+        for sh in reversed(self.shape):
+            index.append(integer_index % sh)
+            integer_index //= sh
+        index.reverse()
+        return tuple(index)
+
+    def _check_symbolic_index(self, index):
+        # Check if any index is symbolic:
+        tuple_index = (index if isinstance(index, tuple) else (index,))
+        if any((isinstance(i, Expr) and (not i.is_number)) for i in tuple_index):
+            for i, nth_dim in zip(tuple_index, self.shape):
+                if ((i < 0) == True) or ((i >= nth_dim) == True):
+                    raise ValueError("index out of range")
+            from sympy.tensor import Indexed
+            return Indexed(self, *tuple_index)
+        return None
+
+    def _setter_iterable_check(self, value):
+        from sympy.matrices.matrixbase import MatrixBase
+        if isinstance(value, (Iterable, MatrixBase, NDimArray)):
+            raise NotImplementedError
+
+    @classmethod
+    def _scan_iterable_shape(cls, iterable):
+        def f(pointer):
+            if not isinstance(pointer, Iterable):
+                return [pointer], ()
+
+            if len(pointer) == 0:
+                return [], (0,)
+
+            result = []
+            elems, shapes = zip(*[f(i) for i in pointer])
+            if len(set(shapes)) != 1:
+                raise ValueError("could not determine shape unambiguously")
+            for i in elems:
+                result.extend(i)
+            return result, (len(shapes),)+shapes[0]
+
+        return f(iterable)
+
+    @classmethod
+    def _handle_ndarray_creation_inputs(cls, iterable=None, shape=None, **kwargs):
+        from sympy.matrices.matrixbase import MatrixBase
+        from sympy.tensor.array import SparseNDimArray
+
+        if shape is None:
+            if iterable is None:
+                shape = ()
+                iterable = ()
+            # Construction of a sparse array from a sparse array
+            elif isinstance(iterable, SparseNDimArray):
+                return iterable._shape, iterable._sparse_array
+
+            # Construct N-dim array from another N-dim array:
+            elif isinstance(iterable, NDimArray):
+                shape = iterable.shape
+
+            # Construct N-dim array from an iterable (numpy arrays included):
+            elif isinstance(iterable, Iterable):
+                iterable, shape = cls._scan_iterable_shape(iterable)
+
+            # Construct N-dim array from a Matrix:
+            elif isinstance(iterable, MatrixBase):
+                shape = iterable.shape
+
+            else:
+                shape = ()
+                iterable = (iterable,)
+
+        if isinstance(iterable, (Dict, dict)) and shape is not None:
+            new_dict = iterable.copy()
+            for k in new_dict:
+                if isinstance(k, (tuple, Tuple)):
+                    new_key = 0
+                    for i, idx in enumerate(k):
+                        new_key = new_key * shape[i] + idx
+                    iterable[new_key] = iterable[k]
+                    del iterable[k]
+
+        if isinstance(shape, (SYMPY_INTS, Integer)):
+            shape = (shape,)
+
+        if not all(isinstance(dim, (SYMPY_INTS, Integer)) for dim in shape):
+            raise TypeError("Shape should contain integers only.")
+
+        return tuple(shape), iterable
+
+    def __len__(self):
+        """Overload common function len(). Returns number of elements in array.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray.zeros(3, 3)
+        >>> a
+        [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
+        >>> len(a)
+        9
+
+        """
+        return self._loop_size
+
+    @property
+    def shape(self):
+        """
+        Returns array shape (dimension).
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray.zeros(3, 3)
+        >>> a.shape
+        (3, 3)
+
+        """
+        return self._shape
+
+    def rank(self):
+        """
+        Returns rank of array.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray.zeros(3,4,5,6,3)
+        >>> a.rank()
+        5
+
+        """
+        return self._rank
+
+    def diff(self, *args, **kwargs):
+        """
+        Calculate the derivative of each element in the array.
+
+        Examples
+        ========
+
+        >>> from sympy import ImmutableDenseNDimArray
+        >>> from sympy.abc import x, y
+        >>> M = ImmutableDenseNDimArray([[x, y], [1, x*y]])
+        >>> M.diff(x)
+        [[1, 0], [0, y]]
+
+        """
+        from sympy.tensor.array.array_derivatives import ArrayDerivative
+        kwargs.setdefault('evaluate', True)
+        return ArrayDerivative(self.as_immutable(), *args, **kwargs)
+
+    def _eval_derivative(self, base):
+        # Types are (base: scalar, self: array)
+        return self.applyfunc(lambda x: base.diff(x))
+
+    def _eval_derivative_n_times(self, s, n):
+        return Basic._eval_derivative_n_times(self, s, n)
+
+    def applyfunc(self, f):
+        """Apply a function to each element of the N-dim array.
+
+        Examples
+        ========
+
+        >>> from sympy import ImmutableDenseNDimArray
+        >>> m = ImmutableDenseNDimArray([i*2+j for i in range(2) for j in range(2)], (2, 2))
+        >>> m
+        [[0, 1], [2, 3]]
+        >>> m.applyfunc(lambda i: 2*i)
+        [[0, 2], [4, 6]]
+        """
+        from sympy.tensor.array import SparseNDimArray
+        from sympy.tensor.array.arrayop import Flatten
+
+        if isinstance(self, SparseNDimArray) and f(S.Zero) == 0:
+            return type(self)({k: f(v) for k, v in self._sparse_array.items() if f(v) != 0}, self.shape)
+
+        return type(self)(map(f, Flatten(self)), self.shape)
+
+    def _sympystr(self, printer):
+        def f(sh, shape_left, i, j):
+            if len(shape_left) == 1:
+                return "["+", ".join([printer._print(self[self._get_tuple_index(e)]) for e in range(i, j)])+"]"
+
+            sh //= shape_left[0]
+            return "[" + ", ".join([f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh) for e in range(shape_left[0])]) + "]" # + "\n"*len(shape_left)
+
+        if self.rank() == 0:
+            return printer._print(self[()])
+        if 0 in self.shape:
+            return f"{self.__class__.__name__}([], {self.shape})"
+        return f(self._loop_size, self.shape, 0, self._loop_size)
+
+    def tolist(self):
+        """
+        Converting MutableDenseNDimArray to one-dim list
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray([1, 2, 3, 4], (2, 2))
+        >>> a
+        [[1, 2], [3, 4]]
+        >>> b = a.tolist()
+        >>> b
+        [[1, 2], [3, 4]]
+        """
+
+        def f(sh, shape_left, i, j):
+            if len(shape_left) == 1:
+                return [self[self._get_tuple_index(e)] for e in range(i, j)]
+            result = []
+            sh //= shape_left[0]
+            for e in range(shape_left[0]):
+                result.append(f(sh, shape_left[1:], i+e*sh, i+(e+1)*sh))
+            return result
+
+        return f(self._loop_size, self.shape, 0, self._loop_size)
+
+    def __add__(self, other):
+        from sympy.tensor.array.arrayop import Flatten
+
+        if not isinstance(other, NDimArray):
+            return NotImplemented
+
+        if self.shape != other.shape:
+            raise ValueError("array shape mismatch")
+        result_list = [i+j for i,j in zip(Flatten(self), Flatten(other))]
+
+        return type(self)(result_list, self.shape)
+
+    def __sub__(self, other):
+        from sympy.tensor.array.arrayop import Flatten
+
+        if not isinstance(other, NDimArray):
+            return NotImplemented
+
+        if self.shape != other.shape:
+            raise ValueError("array shape mismatch")
+        result_list = [i-j for i,j in zip(Flatten(self), Flatten(other))]
+
+        return type(self)(result_list, self.shape)
+
+    def __mul__(self, other):
+        from sympy.matrices.matrixbase import MatrixBase
+        from sympy.tensor.array import SparseNDimArray
+        from sympy.tensor.array.arrayop import Flatten
+
+        if isinstance(other, (Iterable, NDimArray, MatrixBase)):
+            raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
+
+        other = sympify(other)
+        if isinstance(self, SparseNDimArray):
+            if other.is_zero:
+                return type(self)({}, self.shape)
+            return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape)
+
+        result_list = [i*other for i in Flatten(self)]
+        return type(self)(result_list, self.shape)
+
+    def __rmul__(self, other):
+        from sympy.matrices.matrixbase import MatrixBase
+        from sympy.tensor.array import SparseNDimArray
+        from sympy.tensor.array.arrayop import Flatten
+
+        if isinstance(other, (Iterable, NDimArray, MatrixBase)):
+            raise ValueError("scalar expected, use tensorproduct(...) for tensorial product")
+
+        other = sympify(other)
+        if isinstance(self, SparseNDimArray):
+            if other.is_zero:
+                return type(self)({}, self.shape)
+            return type(self)({k: other*v for (k, v) in self._sparse_array.items()}, self.shape)
+
+        result_list = [other*i for i in Flatten(self)]
+        return type(self)(result_list, self.shape)
+
+    def __truediv__(self, other):
+        from sympy.matrices.matrixbase import MatrixBase
+        from sympy.tensor.array import SparseNDimArray
+        from sympy.tensor.array.arrayop import Flatten
+
+        if isinstance(other, (Iterable, NDimArray, MatrixBase)):
+            raise ValueError("scalar expected")
+
+        other = sympify(other)
+        if isinstance(self, SparseNDimArray) and other != S.Zero:
+            return type(self)({k: v/other for (k, v) in self._sparse_array.items()}, self.shape)
+
+        result_list = [i/other for i in Flatten(self)]
+        return type(self)(result_list, self.shape)
+
+    def __rtruediv__(self, other):
+        raise NotImplementedError('unsupported operation on NDimArray')
+
+    def __neg__(self):
+        from sympy.tensor.array import SparseNDimArray
+        from sympy.tensor.array.arrayop import Flatten
+
+        if isinstance(self, SparseNDimArray):
+            return type(self)({k: -v for (k, v) in self._sparse_array.items()}, self.shape)
+
+        result_list = [-i for i in Flatten(self)]
+        return type(self)(result_list, self.shape)
+
+    def __iter__(self):
+        def iterator():
+            if self._shape:
+                for i in range(self._shape[0]):
+                    yield self[i]
+            else:
+                yield self[()]
+
+        return iterator()
+
+    def __eq__(self, other):
+        """
+        NDimArray instances can be compared to each other.
+        Instances equal if they have same shape and data.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableDenseNDimArray
+        >>> a = MutableDenseNDimArray.zeros(2, 3)
+        >>> b = MutableDenseNDimArray.zeros(2, 3)
+        >>> a == b
+        True
+        >>> c = a.reshape(3, 2)
+        >>> c == b
+        False
+        >>> a[0,0] = 1
+        >>> b[0,0] = 2
+        >>> a == b
+        False
+        """
+        from sympy.tensor.array import SparseNDimArray
+        if not isinstance(other, NDimArray):
+            return False
+
+        if not self.shape == other.shape:
+            return False
+
+        if isinstance(self, SparseNDimArray) and isinstance(other, SparseNDimArray):
+            return dict(self._sparse_array) == dict(other._sparse_array)
+
+        return list(self) == list(other)
+
+    def __ne__(self, other):
+        return not self == other
+
+    def _eval_transpose(self):
+        if self.rank() != 2:
+            raise ValueError("array rank not 2")
+        from .arrayop import permutedims
+        return permutedims(self, (1, 0))
+
+    def transpose(self):
+        return self._eval_transpose()
+
+    def _eval_conjugate(self):
+        from sympy.tensor.array.arrayop import Flatten
+
+        return self.func([i.conjugate() for i in Flatten(self)], self.shape)
+
+    def conjugate(self):
+        return self._eval_conjugate()
+
+    def _eval_adjoint(self):
+        return self.transpose().conjugate()
+
+    def adjoint(self):
+        return self._eval_adjoint()
+
+    def _slice_expand(self, s, dim):
+        if not isinstance(s, slice):
+                return (s,)
+        start, stop, step = s.indices(dim)
+        return [start + i*step for i in range((stop-start)//step)]
+
+    def _get_slice_data_for_array_access(self, index):
+        sl_factors = [self._slice_expand(i, dim) for (i, dim) in zip(index, self.shape)]
+        eindices = itertools.product(*sl_factors)
+        return sl_factors, eindices
+
+    def _get_slice_data_for_array_assignment(self, index, value):
+        if not isinstance(value, NDimArray):
+            value = type(self)(value)
+        sl_factors, eindices = self._get_slice_data_for_array_access(index)
+        slice_offsets = [min(i) if isinstance(i, list) else None for i in sl_factors]
+        # TODO: add checks for dimensions for `value`?
+        return value, eindices, slice_offsets
+
+    @classmethod
+    def _check_special_bounds(cls, flat_list, shape):
+        if shape == () and len(flat_list) != 1:
+            raise ValueError("arrays without shape need one scalar value")
+        if shape == (0,) and len(flat_list) > 0:
+            raise ValueError("if array shape is (0,) there cannot be elements")
+
+    def _check_index_for_getitem(self, index):
+        if isinstance(index, (SYMPY_INTS, Integer, slice)):
+            index = (index,)
+
+        if len(index) < self.rank():
+            index = tuple(index) + \
+                          tuple(slice(None) for i in range(len(index), self.rank()))
+
+        if len(index) > self.rank():
+            raise ValueError('Dimension of index greater than rank of array')
+
+        return index
+
+
+class ImmutableNDimArray(NDimArray, Basic):
+    _op_priority = 11.0
+
+    def __hash__(self):
+        return Basic.__hash__(self)
+
+    def as_immutable(self):
+        return self
+
+    def as_mutable(self):
+        raise NotImplementedError("abstract method")

venv/lib/python3.12/site-packages/sympy/tensor/array/sparse_ndim_array.py

@@ -0,0 +1,196 @@
+from sympy.core.basic import Basic
+from sympy.core.containers import (Dict, 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
+from sympy.utilities.iterables import flatten
+
+import functools
+
+class SparseNDimArray(NDimArray):
+
+    def __new__(self, *args, **kwargs):
+        return ImmutableSparseNDimArray(*args, **kwargs)
+
+    def __getitem__(self, index):
+        """
+        Get an element from a sparse N-dim array.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableSparseNDimArray
+        >>> a = MutableSparseNDimArray(range(4), (2, 2))
+        >>> a
+        [[0, 1], [2, 3]]
+        >>> a[0, 0]
+        0
+        >>> a[1, 1]
+        3
+        >>> a[0]
+        [0, 1]
+        >>> a[1]
+        [2, 3]
+
+        Symbolic indexing:
+
+        >>> from sympy.abc import i, j
+        >>> a[i, j]
+        [[0, 1], [2, 3]][i, j]
+
+        Replace `i` and `j` to get element `(0, 0)`:
+
+        >>> a[i, j].subs({i: 0, j: 0})
+        0
+
+        """
+        syindex = self._check_symbolic_index(index)
+        if syindex is not None:
+            return syindex
+
+        index = self._check_index_for_getitem(index)
+
+        # `index` is a tuple with one or more slices:
+        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._sparse_array.get(self._parse_index(i), S.Zero) 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._sparse_array.get(index, S.Zero)
+
+    @classmethod
+    def zeros(cls, *shape):
+        """
+        Return a sparse N-dim array of zeros.
+        """
+        return cls({}, shape)
+
+    def tomatrix(self):
+        """
+        Converts MutableDenseNDimArray to Matrix. Can convert only 2-dim array, else will raise error.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableSparseNDimArray
+        >>> a = MutableSparseNDimArray([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 SparseMatrix
+        if self.rank() != 2:
+            raise ValueError('Dimensions must be of size of 2')
+
+        mat_sparse = {}
+        for key, value in self._sparse_array.items():
+            mat_sparse[self._get_tuple_index(key)] = value
+
+        return SparseMatrix(self.shape[0], self.shape[1], mat_sparse)
+
+    def reshape(self, *newshape):
+        new_total_size = functools.reduce(lambda x,y: x*y, newshape)
+        if new_total_size != self._loop_size:
+            raise ValueError("Invalid reshape parameters " + newshape)
+
+        return type(self)(self._sparse_array, newshape)
+
+class ImmutableSparseNDimArray(SparseNDimArray, ImmutableNDimArray): # type: ignore
+
+    def __new__(cls, iterable=None, shape=None, **kwargs):
+        shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
+        shape = Tuple(*map(_sympify, shape))
+        cls._check_special_bounds(flat_list, shape)
+        loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
+
+        # Sparse array:
+        if isinstance(flat_list, (dict, Dict)):
+            sparse_array = Dict(flat_list)
+        else:
+            sparse_array = {}
+            for i, el in enumerate(flatten(flat_list)):
+                if el != 0:
+                    sparse_array[i] = _sympify(el)
+
+        sparse_array = Dict(sparse_array)
+
+        self = Basic.__new__(cls, sparse_array, shape, **kwargs)
+        self._shape = shape
+        self._rank = len(shape)
+        self._loop_size = loop_size
+        self._sparse_array = sparse_array
+
+        return self
+
+    def __setitem__(self, index, value):
+        raise TypeError("immutable N-dim array")
+
+    def as_mutable(self):
+        return MutableSparseNDimArray(self)
+
+
+class MutableSparseNDimArray(MutableNDimArray, SparseNDimArray):
+
+    def __new__(cls, iterable=None, shape=None, **kwargs):
+        shape, flat_list = cls._handle_ndarray_creation_inputs(iterable, shape, **kwargs)
+        self = object.__new__(cls)
+        self._shape = shape
+        self._rank = len(shape)
+        self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else len(flat_list)
+
+        # Sparse array:
+        if isinstance(flat_list, (dict, Dict)):
+            self._sparse_array = dict(flat_list)
+            return self
+
+        self._sparse_array = {}
+
+        for i, el in enumerate(flatten(flat_list)):
+            if el != 0:
+                self._sparse_array[i] = _sympify(el)
+
+        return self
+
+    def __setitem__(self, index, value):
+        """Allows to set items to MutableDenseNDimArray.
+
+        Examples
+        ========
+
+        >>> from sympy import MutableSparseNDimArray
+        >>> a = MutableSparseNDimArray.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]
+                other_value = value[other_i]
+                complete_index = self._parse_index(i)
+                if other_value != 0:
+                    self._sparse_array[complete_index] = other_value
+                elif complete_index in self._sparse_array:
+                    self._sparse_array.pop(complete_index)
+        else:
+            index = self._parse_index(index)
+            value = _sympify(value)
+            if value == 0 and index in self._sparse_array:
+                self._sparse_array.pop(index)
+            else:
+                self._sparse_array[index] = value
+
+    def as_immutable(self):
+        return ImmutableSparseNDimArray(self)
+
+    @property
+    def free_symbols(self):
+        return {i for j in self._sparse_array.values() for i in j.free_symbols}

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_array_comprehension.py

@@ -0,0 +1,78 @@
+from sympy.tensor.array.array_comprehension import ArrayComprehension, ArrayComprehensionMap
+from sympy.tensor.array import ImmutableDenseNDimArray
+from sympy.abc import i, j, k, l
+from sympy.testing.pytest import raises
+from sympy.matrices import Matrix
+
+
+def test_array_comprehension():
+    a = ArrayComprehension(i*j, (i, 1, 3), (j, 2, 4))
+    b = ArrayComprehension(i, (i, 1, j+1))
+    c = ArrayComprehension(i+j+k+l, (i, 1, 2), (j, 1, 3), (k, 1, 4), (l, 1, 5))
+    d = ArrayComprehension(k, (i, 1, 5))
+    e = ArrayComprehension(i, (j, k+1, k+5))
+    assert a.doit().tolist() == [[2, 3, 4], [4, 6, 8], [6, 9, 12]]
+    assert a.shape == (3, 3)
+    assert a.is_shape_numeric == True
+    assert a.tolist() == [[2, 3, 4], [4, 6, 8], [6, 9, 12]]
+    assert a.tomatrix() == Matrix([
+                           [2, 3, 4],
+                           [4, 6, 8],
+                           [6, 9, 12]])
+    assert len(a) == 9
+    assert isinstance(b.doit(), ArrayComprehension)
+    assert isinstance(a.doit(), ImmutableDenseNDimArray)
+    assert b.subs(j, 3) == ArrayComprehension(i, (i, 1, 4))
+    assert b.free_symbols == {j}
+    assert b.shape == (j + 1,)
+    assert b.rank() == 1
+    assert b.is_shape_numeric == False
+    assert c.free_symbols == set()
+    assert c.function == i + j + k + l
+    assert c.limits == ((i, 1, 2), (j, 1, 3), (k, 1, 4), (l, 1, 5))
+    assert c.doit().tolist() == [[[[4, 5, 6, 7, 8], [5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11]],
+                                  [[5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12]],
+                                  [[6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13]]],
+                                 [[[5, 6, 7, 8, 9], [6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12]],
+                                  [[6, 7, 8, 9, 10], [7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13]],
+                                  [[7, 8, 9, 10, 11], [8, 9, 10, 11, 12], [9, 10, 11, 12, 13], [10, 11, 12, 13, 14]]]]
+    assert c.free_symbols == set()
+    assert c.variables == [i, j, k, l]
+    assert c.bound_symbols == [i, j, k, l]
+    assert d.doit().tolist() == [k, k, k, k, k]
+    assert len(e) == 5
+    raises(TypeError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, [1, 3, 2])))
+    raises(ValueError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, 1)))
+    raises(ValueError, lambda: ArrayComprehension(i*j, (i, 1, 3), (j, 2, j+1)))
+    raises(ValueError, lambda: len(ArrayComprehension(i*j, (i, 1, 3), (j, 2, j+4))))
+    raises(TypeError, lambda: ArrayComprehension(i*j, (i, 0, i + 1.5), (j, 0, 2)))
+    raises(ValueError, lambda: b.tolist())
+    raises(ValueError, lambda: b.tomatrix())
+    raises(ValueError, lambda: c.tomatrix())
+
+def test_arraycomprehensionmap():
+    a = ArrayComprehensionMap(lambda i: i+1, (i, 1, 5))
+    assert a.doit().tolist() == [2, 3, 4, 5, 6]
+    assert a.shape == (5,)
+    assert a.is_shape_numeric
+    assert a.tolist() == [2, 3, 4, 5, 6]
+    assert len(a) == 5
+    assert isinstance(a.doit(), ImmutableDenseNDimArray)
+    expr = ArrayComprehensionMap(lambda i: i+1, (i, 1, k))
+    assert expr.doit() == expr
+    assert expr.subs(k, 4) == ArrayComprehensionMap(lambda i: i+1, (i, 1, 4))
+    assert expr.subs(k, 4).doit() == ImmutableDenseNDimArray([2, 3, 4, 5])
+    b = ArrayComprehensionMap(lambda i: i+1, (i, 1, 2), (i, 1, 3), (i, 1, 4), (i, 1, 5))
+    assert b.doit().tolist() == [[[[2, 3, 4, 5, 6], [3, 5, 7, 9, 11], [4, 7, 10, 13, 16], [5, 9, 13, 17, 21]],
+                                  [[3, 5, 7, 9, 11], [5, 9, 13, 17, 21], [7, 13, 19, 25, 31], [9, 17, 25, 33, 41]],
+                                  [[4, 7, 10, 13, 16], [7, 13, 19, 25, 31], [10, 19, 28, 37, 46], [13, 25, 37, 49, 61]]],
+                                 [[[3, 5, 7, 9, 11], [5, 9, 13, 17, 21], [7, 13, 19, 25, 31], [9, 17, 25, 33, 41]],
+                                  [[5, 9, 13, 17, 21], [9, 17, 25, 33, 41], [13, 25, 37, 49, 61], [17, 33, 49, 65, 81]],
+                                  [[7, 13, 19, 25, 31], [13, 25, 37, 49, 61], [19, 37, 55, 73, 91], [25, 49, 73, 97, 121]]]]
+
+    # tests about lambda expression
+    assert ArrayComprehensionMap(lambda: 3, (i, 1, 5)).doit().tolist() == [3, 3, 3, 3, 3]
+    assert ArrayComprehensionMap(lambda i: i+1, (i, 1, 5)).doit().tolist() == [2, 3, 4, 5, 6]
+    raises(ValueError, lambda: ArrayComprehensionMap(i*j, (i, 1, 3), (j, 2, 4)))
+    a = ArrayComprehensionMap(lambda i, j: i+j, (i, 1, 5))
+    raises(ValueError, lambda: a.doit())

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_array_derivatives.py

@@ -0,0 +1,52 @@
+from sympy.core.symbol import symbols
+from sympy.matrices.dense import Matrix
+from sympy.matrices.expressions.matexpr import MatrixSymbol
+from sympy.tensor.array.ndim_array import NDimArray
+from sympy.matrices.matrixbase import MatrixBase
+from sympy.tensor.array.array_derivatives import ArrayDerivative
+
+x, y, z, t = symbols("x y z t")
+
+m = Matrix([[x, y], [z, t]])
+
+M = MatrixSymbol("M", 3, 2)
+N = MatrixSymbol("N", 4, 3)
+
+
+def test_array_derivative_construction():
+
+    d = ArrayDerivative(x, m, evaluate=False)
+    assert d.shape == (2, 2)
+    expr = d.doit()
+    assert isinstance(expr, MatrixBase)
+    assert expr.shape == (2, 2)
+
+    d = ArrayDerivative(m, m, evaluate=False)
+    assert d.shape == (2, 2, 2, 2)
+    expr = d.doit()
+    assert isinstance(expr, NDimArray)
+    assert expr.shape == (2, 2, 2, 2)
+
+    d = ArrayDerivative(m, x, evaluate=False)
+    assert d.shape == (2, 2)
+    expr = d.doit()
+    assert isinstance(expr, MatrixBase)
+    assert expr.shape == (2, 2)
+
+    d = ArrayDerivative(M, N, evaluate=False)
+    assert d.shape == (4, 3, 3, 2)
+    expr = d.doit()
+    assert isinstance(expr, ArrayDerivative)
+    assert expr.shape == (4, 3, 3, 2)
+
+    d = ArrayDerivative(M, (N, 2), evaluate=False)
+    assert d.shape == (4, 3, 4, 3, 3, 2)
+    expr = d.doit()
+    assert isinstance(expr, ArrayDerivative)
+    assert expr.shape == (4, 3, 4, 3, 3, 2)
+
+    d = ArrayDerivative(M.as_explicit(), (N.as_explicit(), 2), evaluate=False)
+    assert d.doit().shape == (4, 3, 4, 3, 3, 2)
+    expr = d.doit()
+    assert isinstance(expr, NDimArray)
+    assert expr.shape == (4, 3, 4, 3, 3, 2)

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_arrayop.py

@@ -0,0 +1,361 @@
+import itertools
+import random
+
+from sympy.combinatorics import Permutation
+from sympy.combinatorics.permutations import _af_invert
+from sympy.testing.pytest import raises
+
+from sympy.core.function import diff
+from sympy.core.symbol import symbols
+from sympy.functions.elementary.complexes import (adjoint, conjugate, transpose)
+from sympy.functions.elementary.exponential import (exp, log)
+from sympy.functions.elementary.trigonometric import (cos, sin)
+from sympy.tensor.array import Array, ImmutableDenseNDimArray, ImmutableSparseNDimArray, MutableSparseNDimArray
+
+from sympy.tensor.array.arrayop import tensorproduct, tensorcontraction, derive_by_array, permutedims, Flatten, \
+    tensordiagonal
+
+
+def test_import_NDimArray():
+    from sympy.tensor.array import NDimArray
+    del NDimArray
+
+
+def test_tensorproduct():
+    x,y,z,t = symbols('x y z t')
+    from sympy.abc import a,b,c,d
+    assert tensorproduct() == 1
+    assert tensorproduct([x]) == Array([x])
+    assert tensorproduct([x], [y]) == Array([[x*y]])
+    assert tensorproduct([x], [y], [z]) == Array([[[x*y*z]]])
+    assert tensorproduct([x], [y], [z], [t]) == Array([[[[x*y*z*t]]]])
+
+    assert tensorproduct(x) == x
+    assert tensorproduct(x, y) == x*y
+    assert tensorproduct(x, y, z) == x*y*z
+    assert tensorproduct(x, y, z, t) == x*y*z*t
+
+    for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
+        A = ArrayType([x, y])
+        B = ArrayType([1, 2, 3])
+        C = ArrayType([a, b, c, d])
+
+        assert tensorproduct(A, B, C) == ArrayType([[[a*x, b*x, c*x, d*x], [2*a*x, 2*b*x, 2*c*x, 2*d*x], [3*a*x, 3*b*x, 3*c*x, 3*d*x]],
+                                                    [[a*y, b*y, c*y, d*y], [2*a*y, 2*b*y, 2*c*y, 2*d*y], [3*a*y, 3*b*y, 3*c*y, 3*d*y]]])
+
+        assert tensorproduct([x, y], [1, 2, 3]) == tensorproduct(A, B)
+
+        assert tensorproduct(A, 2) == ArrayType([2*x, 2*y])
+        assert tensorproduct(A, [2]) == ArrayType([[2*x], [2*y]])
+        assert tensorproduct([2], A) == ArrayType([[2*x, 2*y]])
+        assert tensorproduct(a, A) == ArrayType([a*x, a*y])
+        assert tensorproduct(a, A, B) == ArrayType([[a*x, 2*a*x, 3*a*x], [a*y, 2*a*y, 3*a*y]])
+        assert tensorproduct(A, B, a) == ArrayType([[a*x, 2*a*x, 3*a*x], [a*y, 2*a*y, 3*a*y]])
+        assert tensorproduct(B, a, A) == ArrayType([[a*x, a*y], [2*a*x, 2*a*y], [3*a*x, 3*a*y]])
+
+    # tests for large scale sparse array
+    for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
+        a = SparseArrayType({1:2, 3:4},(1000, 2000))
+        b = SparseArrayType({1:2, 3:4},(1000, 2000))
+        assert tensorproduct(a, b) == ImmutableSparseNDimArray({2000001: 4, 2000003: 8, 6000001: 8, 6000003: 16}, (1000, 2000, 1000, 2000))
+
+
+def test_tensorcontraction():
+    from sympy.abc import a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x
+    B = Array(range(18), (2, 3, 3))
+    assert tensorcontraction(B, (1, 2)) == Array([12, 39])
+    C1 = Array([a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x], (2, 3, 2, 2))
+
+    assert tensorcontraction(C1, (0, 2)) == Array([[a + o, b + p], [e + s, f + t], [i + w, j + x]])
+    assert tensorcontraction(C1, (0, 2, 3)) == Array([a + p, e + t, i + x])
+    assert tensorcontraction(C1, (2, 3)) == Array([[a + d, e + h, i + l], [m + p, q + t, u + x]])
+
+
+def test_derivative_by_array():
+    from sympy.abc import i, j, t, x, y, z
+
+    bexpr = x*y**2*exp(z)*log(t)
+    sexpr = sin(bexpr)
+    cexpr = cos(bexpr)
+
+    a = Array([sexpr])
+
+    assert derive_by_array(sexpr, t) == x*y**2*exp(z)*cos(x*y**2*exp(z)*log(t))/t
+    assert derive_by_array(sexpr, [x, y, z]) == Array([bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr, bexpr*cexpr])
+    assert derive_by_array(a, [x, y, z]) == Array([[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr], [bexpr*cexpr]])
+
+    assert derive_by_array(sexpr, [[x, y], [z, t]]) == Array([[bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr], [bexpr*cexpr, bexpr/log(t)/t*cexpr]])
+    assert derive_by_array(a, [[x, y], [z, t]]) == Array([[[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr]], [[bexpr*cexpr], [bexpr/log(t)/t*cexpr]]])
+    assert derive_by_array([[x, y], [z, t]], [x, y]) == Array([[[1, 0], [0, 0]], [[0, 1], [0, 0]]])
+    assert derive_by_array([[x, y], [z, t]], [[x, y], [z, t]]) == Array([[[[1, 0], [0, 0]], [[0, 1], [0, 0]]],
+                                                                         [[[0, 0], [1, 0]], [[0, 0], [0, 1]]]])
+
+    assert diff(sexpr, t) == x*y**2*exp(z)*cos(x*y**2*exp(z)*log(t))/t
+    assert diff(sexpr, Array([x, y, z])) == Array([bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr, bexpr*cexpr])
+    assert diff(a, Array([x, y, z])) == Array([[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr], [bexpr*cexpr]])
+
+    assert diff(sexpr, Array([[x, y], [z, t]])) == Array([[bexpr/x*cexpr, 2*y*bexpr/y**2*cexpr], [bexpr*cexpr, bexpr/log(t)/t*cexpr]])
+    assert diff(a, Array([[x, y], [z, t]])) == Array([[[bexpr/x*cexpr], [2*y*bexpr/y**2*cexpr]], [[bexpr*cexpr], [bexpr/log(t)/t*cexpr]]])
+    assert diff(Array([[x, y], [z, t]]), Array([x, y])) == Array([[[1, 0], [0, 0]], [[0, 1], [0, 0]]])
+    assert diff(Array([[x, y], [z, t]]), Array([[x, y], [z, t]])) == Array([[[[1, 0], [0, 0]], [[0, 1], [0, 0]]],
+                                                                         [[[0, 0], [1, 0]], [[0, 0], [0, 1]]]])
+
+    # test for large scale sparse array
+    for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
+        b = MutableSparseNDimArray({0:i, 1:j}, (10000, 20000))
+        assert derive_by_array(b, i) == ImmutableSparseNDimArray({0: 1}, (10000, 20000))
+        assert derive_by_array(b, (i, j)) == ImmutableSparseNDimArray({0: 1, 200000001: 1}, (2, 10000, 20000))
+
+    #https://github.com/sympy/sympy/issues/20655
+    U = Array([x, y, z])
+    E = 2
+    assert derive_by_array(E, U) ==  ImmutableDenseNDimArray([0, 0, 0])
+
+
+def test_issue_emerged_while_discussing_10972():
+    ua = Array([-1,0])
+    Fa = Array([[0, 1], [-1, 0]])
+    po = tensorproduct(Fa, ua, Fa, ua)
+    assert tensorcontraction(po, (1, 2), (4, 5)) == Array([[0, 0], [0, 1]])
+
+    sa = symbols('a0:144')
+    po = Array(sa, [2, 2, 3, 3, 2, 2])
+    assert tensorcontraction(po, (0, 1), (2, 3), (4, 5)) == sa[0] + sa[108] + sa[111] + sa[124] + sa[127] + sa[140] + sa[143] + sa[16] + sa[19] + sa[3] + sa[32] + sa[35]
+    assert tensorcontraction(po, (0, 1, 4, 5), (2, 3)) == sa[0] + sa[111] + sa[127] + sa[143] + sa[16] + sa[32]
+    assert tensorcontraction(po, (0, 1), (4, 5)) == Array([[sa[0] + sa[108] + sa[111] + sa[3], sa[112] + sa[115] + sa[4] + sa[7],
+                                                             sa[11] + sa[116] + sa[119] + sa[8]], [sa[12] + sa[120] + sa[123] + sa[15],
+                                                             sa[124] + sa[127] + sa[16] + sa[19], sa[128] + sa[131] + sa[20] + sa[23]],
+                                                            [sa[132] + sa[135] + sa[24] + sa[27], sa[136] + sa[139] + sa[28] + sa[31],
+                                                             sa[140] + sa[143] + sa[32] + sa[35]]])
+    assert tensorcontraction(po, (0, 1), (2, 3)) == Array([[sa[0] + sa[108] + sa[124] + sa[140] + sa[16] + sa[32], sa[1] + sa[109] + sa[125] + sa[141] + sa[17] + sa[33]],
+                                                           [sa[110] + sa[126] + sa[142] + sa[18] + sa[2] + sa[34], sa[111] + sa[127] + sa[143] + sa[19] + sa[3] + sa[35]]])
+
+
+def test_array_permutedims():
+    sa = symbols('a0:144')
+
+    for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
+        m1 = ArrayType(sa[:6], (2, 3))
+        assert permutedims(m1, (1, 0)) == transpose(m1)
+        assert m1.tomatrix().T == permutedims(m1, (1, 0)).tomatrix()
+
+        assert m1.tomatrix().T == transpose(m1).tomatrix()
+        assert m1.tomatrix().C == conjugate(m1).tomatrix()
+        assert m1.tomatrix().H == adjoint(m1).tomatrix()
+
+        assert m1.tomatrix().T == m1.transpose().tomatrix()
+        assert m1.tomatrix().C == m1.conjugate().tomatrix()
+        assert m1.tomatrix().H == m1.adjoint().tomatrix()
+
+        raises(ValueError, lambda: permutedims(m1, (0,)))
+        raises(ValueError, lambda: permutedims(m1, (0, 0)))
+        raises(ValueError, lambda: permutedims(m1, (1, 2, 0)))
+
+        # Some tests with random arrays:
+        dims = 6
+        shape = [random.randint(1,5) for i in range(dims)]
+        elems = [random.random() for i in range(tensorproduct(*shape))]
+        ra = ArrayType(elems, shape)
+        perm = list(range(dims))
+        # Randomize the permutation:
+        random.shuffle(perm)
+        # Test inverse permutation:
+        assert permutedims(permutedims(ra, perm), _af_invert(perm)) == ra
+        # Test that permuted shape corresponds to action by `Permutation`:
+        assert permutedims(ra, perm).shape == tuple(Permutation(perm)(shape))
+
+        z = ArrayType.zeros(4,5,6,7)
+
+        assert permutedims(z, (2, 3, 1, 0)).shape == (6, 7, 5, 4)
+        assert permutedims(z, [2, 3, 1, 0]).shape == (6, 7, 5, 4)
+        assert permutedims(z, Permutation([2, 3, 1, 0])).shape == (6, 7, 5, 4)
+
+        po = ArrayType(sa, [2, 2, 3, 3, 2, 2])
+
+        raises(ValueError, lambda: permutedims(po, (1, 1)))
+        raises(ValueError, lambda: po.transpose())
+        raises(ValueError, lambda: po.adjoint())
+
+        assert permutedims(po, reversed(range(po.rank()))) == ArrayType(
+            [[[[[[sa[0], sa[72]], [sa[36], sa[108]]], [[sa[12], sa[84]], [sa[48], sa[120]]], [[sa[24],
+                                                                                               sa[96]], [sa[60], sa[132]]]],
+               [[[sa[4], sa[76]], [sa[40], sa[112]]], [[sa[16],
+                                                        sa[88]], [sa[52], sa[124]]],
+                [[sa[28], sa[100]], [sa[64], sa[136]]]],
+               [[[sa[8],
+                  sa[80]], [sa[44], sa[116]]], [[sa[20], sa[92]], [sa[56], sa[128]]], [[sa[32],
+                                                                                        sa[104]], [sa[68], sa[140]]]]],
+              [[[[sa[2], sa[74]], [sa[38], sa[110]]], [[sa[14],
+                                                        sa[86]], [sa[50], sa[122]]], [[sa[26], sa[98]], [sa[62], sa[134]]]],
+               [[[sa[6],
+                  sa[78]], [sa[42], sa[114]]], [[sa[18], sa[90]], [sa[54], sa[126]]], [[sa[30],
+                                                                                        sa[102]], [sa[66], sa[138]]]],
+               [[[sa[10], sa[82]], [sa[46], sa[118]]], [[sa[22],
+                                                         sa[94]], [sa[58], sa[130]]],
+                [[sa[34], sa[106]], [sa[70], sa[142]]]]]],
+             [[[[[sa[1],
+                  sa[73]], [sa[37], sa[109]]], [[sa[13], sa[85]], [sa[49], sa[121]]], [[sa[25],
+                                                                                        sa[97]], [sa[61], sa[133]]]],
+               [[[sa[5], sa[77]], [sa[41], sa[113]]], [[sa[17],
+                                                        sa[89]], [sa[53], sa[125]]],
+                [[sa[29], sa[101]], [sa[65], sa[137]]]],
+               [[[sa[9],
+                  sa[81]], [sa[45], sa[117]]], [[sa[21], sa[93]], [sa[57], sa[129]]], [[sa[33],
+                                                                                        sa[105]], [sa[69], sa[141]]]]],
+              [[[[sa[3], sa[75]], [sa[39], sa[111]]], [[sa[15],
+                                                        sa[87]], [sa[51], sa[123]]], [[sa[27], sa[99]], [sa[63], sa[135]]]],
+               [[[sa[7],
+                  sa[79]], [sa[43], sa[115]]], [[sa[19], sa[91]], [sa[55], sa[127]]], [[sa[31],
+                                                                                        sa[103]], [sa[67], sa[139]]]],
+               [[[sa[11], sa[83]], [sa[47], sa[119]]], [[sa[23],
+                                                         sa[95]], [sa[59], sa[131]]],
+                [[sa[35], sa[107]], [sa[71], sa[143]]]]]]])
+
+        assert permutedims(po, (1, 0, 2, 3, 4, 5)) == ArrayType(
+            [[[[[[sa[0], sa[1]], [sa[2], sa[3]]], [[sa[4], sa[5]], [sa[6], sa[7]]], [[sa[8], sa[9]], [sa[10],
+                                                                                                      sa[11]]]],
+               [[[sa[12], sa[13]], [sa[14], sa[15]]], [[sa[16], sa[17]], [sa[18],
+                                                                          sa[19]]], [[sa[20], sa[21]], [sa[22], sa[23]]]],
+               [[[sa[24], sa[25]], [sa[26],
+                                    sa[27]]], [[sa[28], sa[29]], [sa[30], sa[31]]], [[sa[32], sa[33]], [sa[34],
+                                                                                                        sa[35]]]]],
+              [[[[sa[72], sa[73]], [sa[74], sa[75]]], [[sa[76], sa[77]], [sa[78],
+                                                                          sa[79]]], [[sa[80], sa[81]], [sa[82], sa[83]]]],
+               [[[sa[84], sa[85]], [sa[86],
+                                    sa[87]]], [[sa[88], sa[89]], [sa[90], sa[91]]], [[sa[92], sa[93]], [sa[94],
+                                                                                                        sa[95]]]],
+               [[[sa[96], sa[97]], [sa[98], sa[99]]], [[sa[100], sa[101]], [sa[102],
+                                                                            sa[103]]],
+                [[sa[104], sa[105]], [sa[106], sa[107]]]]]], [[[[[sa[36], sa[37]], [sa[38],
+                                                                                    sa[39]]],
+                                                                [[sa[40], sa[41]], [sa[42], sa[43]]],
+                                                                [[sa[44], sa[45]], [sa[46],
+                                                                                    sa[47]]]],
+                                                               [[[sa[48], sa[49]], [sa[50], sa[51]]],
+                                                                [[sa[52], sa[53]], [sa[54],
+                                                                                    sa[55]]],
+                                                                [[sa[56], sa[57]], [sa[58], sa[59]]]],
+                                                               [[[sa[60], sa[61]], [sa[62],
+                                                                                    sa[63]]],
+                                                                [[sa[64], sa[65]], [sa[66], sa[67]]],
+                                                                [[sa[68], sa[69]], [sa[70],
+                                                                                    sa[71]]]]], [
+                                                                  [[[sa[108], sa[109]], [sa[110], sa[111]]],
+                                                                   [[sa[112], sa[113]], [sa[114],
+                                                                                         sa[115]]],
+                                                                   [[sa[116], sa[117]], [sa[118], sa[119]]]],
+                                                                  [[[sa[120], sa[121]], [sa[122],
+                                                                                         sa[123]]],
+                                                                   [[sa[124], sa[125]], [sa[126], sa[127]]],
+                                                                   [[sa[128], sa[129]], [sa[130],
+                                                                                         sa[131]]]],
+                                                                  [[[sa[132], sa[133]], [sa[134], sa[135]]],
+                                                                   [[sa[136], sa[137]], [sa[138],
+                                                                                         sa[139]]],
+                                                                   [[sa[140], sa[141]], [sa[142], sa[143]]]]]]])
+
+        assert permutedims(po, (0, 2, 1, 4, 3, 5)) == ArrayType(
+            [[[[[[sa[0], sa[1]], [sa[4], sa[5]], [sa[8], sa[9]]], [[sa[2], sa[3]], [sa[6], sa[7]], [sa[10],
+                                                                                                    sa[11]]]],
+               [[[sa[36], sa[37]], [sa[40], sa[41]], [sa[44], sa[45]]], [[sa[38],
+                                                                          sa[39]], [sa[42], sa[43]], [sa[46], sa[47]]]]],
+              [[[[sa[12], sa[13]], [sa[16],
+                                    sa[17]], [sa[20], sa[21]]], [[sa[14], sa[15]], [sa[18], sa[19]], [sa[22],
+                                                                                                      sa[23]]]],
+               [[[sa[48], sa[49]], [sa[52], sa[53]], [sa[56], sa[57]]], [[sa[50],
+                                                                          sa[51]], [sa[54], sa[55]], [sa[58], sa[59]]]]],
+              [[[[sa[24], sa[25]], [sa[28],
+                                    sa[29]], [sa[32], sa[33]]], [[sa[26], sa[27]], [sa[30], sa[31]], [sa[34],
+                                                                                                      sa[35]]]],
+               [[[sa[60], sa[61]], [sa[64], sa[65]], [sa[68], sa[69]]], [[sa[62],
+                                                                          sa[63]], [sa[66], sa[67]], [sa[70], sa[71]]]]]],
+             [[[[[sa[72], sa[73]], [sa[76],
+                                    sa[77]], [sa[80], sa[81]]], [[sa[74], sa[75]], [sa[78], sa[79]], [sa[82],
+                                                                                                      sa[83]]]],
+               [[[sa[108], sa[109]], [sa[112], sa[113]], [sa[116], sa[117]]], [[sa[110],
+                                                                                sa[111]], [sa[114], sa[115]],
+                                                                               [sa[118], sa[119]]]]],
+              [[[[sa[84], sa[85]], [sa[88],
+                                    sa[89]], [sa[92], sa[93]]], [[sa[86], sa[87]], [sa[90], sa[91]], [sa[94],
+                                                                                                      sa[95]]]],
+               [[[sa[120], sa[121]], [sa[124], sa[125]], [sa[128], sa[129]]], [[sa[122],
+                                                                                sa[123]], [sa[126], sa[127]],
+                                                                               [sa[130], sa[131]]]]],
+              [[[[sa[96], sa[97]], [sa[100],
+                                    sa[101]], [sa[104], sa[105]]], [[sa[98], sa[99]], [sa[102], sa[103]], [sa[106],
+                                                                                                           sa[107]]]],
+               [[[sa[132], sa[133]], [sa[136], sa[137]], [sa[140], sa[141]]], [[sa[134],
+                                                                                sa[135]], [sa[138], sa[139]],
+                                                                               [sa[142], sa[143]]]]]]])
+
+        po2 = po.reshape(4, 9, 2, 2)
+        assert po2 == ArrayType([[[[sa[0], sa[1]], [sa[2], sa[3]]], [[sa[4], sa[5]], [sa[6], sa[7]]], [[sa[8], sa[9]], [sa[10], sa[11]]], [[sa[12], sa[13]], [sa[14], sa[15]]], [[sa[16], sa[17]], [sa[18], sa[19]]], [[sa[20], sa[21]], [sa[22], sa[23]]], [[sa[24], sa[25]], [sa[26], sa[27]]], [[sa[28], sa[29]], [sa[30], sa[31]]], [[sa[32], sa[33]], [sa[34], sa[35]]]], [[[sa[36], sa[37]], [sa[38], sa[39]]], [[sa[40], sa[41]], [sa[42], sa[43]]], [[sa[44], sa[45]], [sa[46], sa[47]]], [[sa[48], sa[49]], [sa[50], sa[51]]], [[sa[52], sa[53]], [sa[54], sa[55]]], [[sa[56], sa[57]], [sa[58], sa[59]]], [[sa[60], sa[61]], [sa[62], sa[63]]], [[sa[64], sa[65]], [sa[66], sa[67]]], [[sa[68], sa[69]], [sa[70], sa[71]]]], [[[sa[72], sa[73]], [sa[74], sa[75]]], [[sa[76], sa[77]], [sa[78], sa[79]]], [[sa[80], sa[81]], [sa[82], sa[83]]], [[sa[84], sa[85]], [sa[86], sa[87]]], [[sa[88], sa[89]], [sa[90], sa[91]]], [[sa[92], sa[93]], [sa[94], sa[95]]], [[sa[96], sa[97]], [sa[98], sa[99]]], [[sa[100], sa[101]], [sa[102], sa[103]]], [[sa[104], sa[105]], [sa[106], sa[107]]]], [[[sa[108], sa[109]], [sa[110], sa[111]]], [[sa[112], sa[113]], [sa[114], sa[115]]], [[sa[116], sa[117]], [sa[118], sa[119]]], [[sa[120], sa[121]], [sa[122], sa[123]]], [[sa[124], sa[125]], [sa[126], sa[127]]], [[sa[128], sa[129]], [sa[130], sa[131]]], [[sa[132], sa[133]], [sa[134], sa[135]]], [[sa[136], sa[137]], [sa[138], sa[139]]], [[sa[140], sa[141]], [sa[142], sa[143]]]]])
+
+        assert permutedims(po2, (3, 2, 0, 1)) == ArrayType([[[[sa[0], sa[4], sa[8], sa[12], sa[16], sa[20], sa[24], sa[28], sa[32]], [sa[36], sa[40], sa[44], sa[48], sa[52], sa[56], sa[60], sa[64], sa[68]], [sa[72], sa[76], sa[80], sa[84], sa[88], sa[92], sa[96], sa[100], sa[104]], [sa[108], sa[112], sa[116], sa[120], sa[124], sa[128], sa[132], sa[136], sa[140]]], [[sa[2], sa[6], sa[10], sa[14], sa[18], sa[22], sa[26], sa[30], sa[34]], [sa[38], sa[42], sa[46], sa[50], sa[54], sa[58], sa[62], sa[66], sa[70]], [sa[74], sa[78], sa[82], sa[86], sa[90], sa[94], sa[98], sa[102], sa[106]], [sa[110], sa[114], sa[118], sa[122], sa[126], sa[130], sa[134], sa[138], sa[142]]]], [[[sa[1], sa[5], sa[9], sa[13], sa[17], sa[21], sa[25], sa[29], sa[33]], [sa[37], sa[41], sa[45], sa[49], sa[53], sa[57], sa[61], sa[65], sa[69]], [sa[73], sa[77], sa[81], sa[85], sa[89], sa[93], sa[97], sa[101], sa[105]], [sa[109], sa[113], sa[117], sa[121], sa[125], sa[129], sa[133], sa[137], sa[141]]], [[sa[3], sa[7], sa[11], sa[15], sa[19], sa[23], sa[27], sa[31], sa[35]], [sa[39], sa[43], sa[47], sa[51], sa[55], sa[59], sa[63], sa[67], sa[71]], [sa[75], sa[79], sa[83], sa[87], sa[91], sa[95], sa[99], sa[103], sa[107]], [sa[111], sa[115], sa[119], sa[123], sa[127], sa[131], sa[135], sa[139], sa[143]]]]])
+
+    # test for large scale sparse array
+    for SparseArrayType in [ImmutableSparseNDimArray, MutableSparseNDimArray]:
+        A = SparseArrayType({1:1, 10000:2}, (10000, 20000, 10000))
+        assert permutedims(A, (0, 1, 2)) == A
+        assert permutedims(A, (1, 0, 2)) == SparseArrayType({1: 1, 100000000: 2}, (20000, 10000, 10000))
+        B = SparseArrayType({1:1, 20000:2}, (10000, 20000))
+        assert B.transpose() == SparseArrayType({10000: 1, 1: 2}, (20000, 10000))
+
+
+def test_permutedims_with_indices():
+    A = Array(range(32)).reshape(2, 2, 2, 2, 2)
+    indices_new = list("abcde")
+    indices_old = list("ebdac")
+    new_A = permutedims(A, index_order_new=indices_new, index_order_old=indices_old)
+    for a, b, c, d, e in itertools.product(range(2), range(2), range(2), range(2), range(2)):
+        assert new_A[a, b, c, d, e] == A[e, b, d, a, c]
+    indices_old = list("cabed")
+    new_A = permutedims(A, index_order_new=indices_new, index_order_old=indices_old)
+    for a, b, c, d, e in itertools.product(range(2), range(2), range(2), range(2), range(2)):
+        assert new_A[a, b, c, d, e] == A[c, a, b, e, d]
+    raises(ValueError, lambda: permutedims(A, index_order_old=list("aacde"), index_order_new=list("abcde")))
+    raises(ValueError, lambda: permutedims(A, index_order_old=list("abcde"), index_order_new=list("abcce")))
+    raises(ValueError, lambda: permutedims(A, index_order_old=list("abcde"), index_order_new=list("abce")))
+    raises(ValueError, lambda: permutedims(A, index_order_old=list("abce"), index_order_new=list("abce")))
+    raises(ValueError, lambda: permutedims(A, [2, 1, 0, 3, 4], index_order_old=list("abcde")))
+    raises(ValueError, lambda: permutedims(A, [2, 1, 0, 3, 4], index_order_new=list("abcde")))
+
+
+def test_flatten():
+    from sympy.matrices.dense import Matrix
+    for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray, Matrix]:
+        A = ArrayType(range(24)).reshape(4, 6)
+        assert list(Flatten(A)) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
+
+        for i, v in enumerate(Flatten(A)):
+            assert i == v
+
+
+def test_tensordiagonal():
+    from sympy.matrices.dense import eye
+    expr = Array(range(9)).reshape(3, 3)
+    raises(ValueError, lambda: tensordiagonal(expr, [0], [1]))
+    raises(ValueError, lambda: tensordiagonal(expr, [0, 0]))
+    assert tensordiagonal(eye(3), [0, 1]) == Array([1, 1, 1])
+    assert tensordiagonal(expr, [0, 1]) == Array([0, 4, 8])
+    x, y, z = symbols("x y z")
+    expr2 = tensorproduct([x, y, z], expr)
+    assert tensordiagonal(expr2, [1, 2]) == Array([[0, 4*x, 8*x], [0, 4*y, 8*y], [0, 4*z, 8*z]])
+    assert tensordiagonal(expr2, [0, 1]) == Array([[0, 3*y, 6*z], [x, 4*y, 7*z], [2*x, 5*y, 8*z]])
+    assert tensordiagonal(expr2, [0, 1, 2]) == Array([0, 4*y, 8*z])
+    # assert tensordiagonal(expr2, [0]) == permutedims(expr2, [1, 2, 0])
+    # assert tensordiagonal(expr2, [1]) == permutedims(expr2, [0, 2, 1])
+    # assert tensordiagonal(expr2, [2]) == expr2
+    # assert tensordiagonal(expr2, [1], [2]) == expr2
+    # assert tensordiagonal(expr2, [0], [1]) == permutedims(expr2, [2, 0, 1])
+
+    a, b, c, X, Y, Z = symbols("a b c X Y Z")
+    expr3 = tensorproduct([x, y, z], [1, 2, 3], [a, b, c], [X, Y, Z])
+    assert tensordiagonal(expr3, [0, 1, 2, 3]) == Array([x*a*X, 2*y*b*Y, 3*z*c*Z])
+    assert tensordiagonal(expr3, [0, 1], [2, 3]) == tensorproduct([x, 2*y, 3*z], [a*X, b*Y, c*Z])
+
+    # assert tensordiagonal(expr3, [0], [1, 2], [3]) == tensorproduct([x, y, z], [a, 2*b, 3*c], [X, Y, Z])
+    assert tensordiagonal(tensordiagonal(expr3, [2, 3]), [0, 1]) == tensorproduct([a*X, b*Y, c*Z], [x, 2*y, 3*z])
+
+    raises(ValueError, lambda: tensordiagonal([[1, 2, 3], [4, 5, 6]], [0, 1]))
+    raises(ValueError, lambda: tensordiagonal(expr3.reshape(3, 3, 9), [1, 2]))

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_immutable_ndim_array.py

@@ -0,0 +1,452 @@
+from copy import copy
+
+from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
+from sympy.core.containers import Dict
+from sympy.core.function import diff
+from sympy.core.numbers import Rational
+from sympy.core.singleton import S
+from sympy.core.symbol import (Symbol, symbols)
+from sympy.matrices import SparseMatrix
+from sympy.tensor.indexed import (Indexed, IndexedBase)
+from sympy.matrices import Matrix
+from sympy.tensor.array.sparse_ndim_array import ImmutableSparseNDimArray
+from sympy.testing.pytest import raises
+
+
+def test_ndim_array_initiation():
+    arr_with_no_elements = ImmutableDenseNDimArray([], shape=(0,))
+    assert len(arr_with_no_elements) == 0
+    assert arr_with_no_elements.rank() == 1
+
+    raises(ValueError, lambda: ImmutableDenseNDimArray([0], shape=(0,)))
+    raises(ValueError, lambda: ImmutableDenseNDimArray([1, 2, 3], shape=(0,)))
+    raises(ValueError, lambda: ImmutableDenseNDimArray([], shape=()))
+
+    raises(ValueError, lambda: ImmutableSparseNDimArray([0], shape=(0,)))
+    raises(ValueError, lambda: ImmutableSparseNDimArray([1, 2, 3], shape=(0,)))
+    raises(ValueError, lambda: ImmutableSparseNDimArray([], shape=()))
+
+    arr_with_one_element = ImmutableDenseNDimArray([23])
+    assert len(arr_with_one_element) == 1
+    assert arr_with_one_element[0] == 23
+    assert arr_with_one_element[:] == ImmutableDenseNDimArray([23])
+    assert arr_with_one_element.rank() == 1
+
+    arr_with_symbol_element = ImmutableDenseNDimArray([Symbol('x')])
+    assert len(arr_with_symbol_element) == 1
+    assert arr_with_symbol_element[0] == Symbol('x')
+    assert arr_with_symbol_element[:] == ImmutableDenseNDimArray([Symbol('x')])
+    assert arr_with_symbol_element.rank() == 1
+
+    number5 = 5
+    vector = ImmutableDenseNDimArray.zeros(number5)
+    assert len(vector) == number5
+    assert vector.shape == (number5,)
+    assert vector.rank() == 1
+
+    vector = ImmutableSparseNDimArray.zeros(number5)
+    assert len(vector) == number5
+    assert vector.shape == (number5,)
+    assert vector._sparse_array == Dict()
+    assert vector.rank() == 1
+
+    n_dim_array = ImmutableDenseNDimArray(range(3**4), (3, 3, 3, 3,))
+    assert len(n_dim_array) == 3 * 3 * 3 * 3
+    assert n_dim_array.shape == (3, 3, 3, 3)
+    assert n_dim_array.rank() == 4
+
+    array_shape = (3, 3, 3, 3)
+    sparse_array = ImmutableSparseNDimArray.zeros(*array_shape)
+    assert len(sparse_array._sparse_array) == 0
+    assert len(sparse_array) == 3 * 3 * 3 * 3
+    assert n_dim_array.shape == array_shape
+    assert n_dim_array.rank() == 4
+
+    one_dim_array = ImmutableDenseNDimArray([2, 3, 1])
+    assert len(one_dim_array) == 3
+    assert one_dim_array.shape == (3,)
+    assert one_dim_array.rank() == 1
+    assert one_dim_array.tolist() == [2, 3, 1]
+
+    shape = (3, 3)
+    array_with_many_args = ImmutableSparseNDimArray.zeros(*shape)
+    assert len(array_with_many_args) == 3 * 3
+    assert array_with_many_args.shape == shape
+    assert array_with_many_args[0, 0] == 0
+    assert array_with_many_args.rank() == 2
+
+    shape = (int(3), int(3))
+    array_with_long_shape = ImmutableSparseNDimArray.zeros(*shape)
+    assert len(array_with_long_shape) == 3 * 3
+    assert array_with_long_shape.shape == shape
+    assert array_with_long_shape[int(0), int(0)] == 0
+    assert array_with_long_shape.rank() == 2
+
+    vector_with_long_shape = ImmutableDenseNDimArray(range(5), int(5))
+    assert len(vector_with_long_shape) == 5
+    assert vector_with_long_shape.shape == (int(5),)
+    assert vector_with_long_shape.rank() == 1
+    raises(ValueError, lambda: vector_with_long_shape[int(5)])
+
+    from sympy.abc import x
+    for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
+        rank_zero_array = ArrayType(x)
+        assert len(rank_zero_array) == 1
+        assert rank_zero_array.shape == ()
+        assert rank_zero_array.rank() == 0
+        assert rank_zero_array[()] == x
+        raises(ValueError, lambda: rank_zero_array[0])
+
+
+def test_reshape():
+    array = ImmutableDenseNDimArray(range(50), 50)
+    assert array.shape == (50,)
+    assert array.rank() == 1
+
+    array = array.reshape(5, 5, 2)
+    assert array.shape == (5, 5, 2)
+    assert array.rank() == 3
+    assert len(array) == 50
+
+
+def test_getitem():
+    for ArrayType in [ImmutableDenseNDimArray, ImmutableSparseNDimArray]:
+        array = ArrayType(range(24)).reshape(2, 3, 4)
+        assert array.tolist() == [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]]
+        assert array[0] == ArrayType([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
+        assert array[0, 0] == ArrayType([0, 1, 2, 3])
+        value = 0
+        for i in range(2):
+            for j in range(3):
+                for k in range(4):
+                    assert array[i, j, k] == value
+                    value += 1
+
+    raises(ValueError, lambda: array[3, 4, 5])
+    raises(ValueError, lambda: array[3, 4, 5, 6])
+    raises(ValueError, lambda: array[3, 4, 5, 3:4])
+
+
+def test_iterator():
+    array = ImmutableDenseNDimArray(range(4), (2, 2))
+    assert array[0] == ImmutableDenseNDimArray([0, 1])
+    assert array[1] == ImmutableDenseNDimArray([2, 3])
+
+    array = array.reshape(4)
+    j = 0
+    for i in array:
+        assert i == j
+        j += 1
+
+
+def test_sparse():
+    sparse_array = ImmutableSparseNDimArray([0, 0, 0, 1], (2, 2))
+    assert len(sparse_array) == 2 * 2
+    # dictionary where all data is, only non-zero entries are actually stored:
+    assert len(sparse_array._sparse_array) == 1
+
+    assert sparse_array.tolist() == [[0, 0], [0, 1]]
+
+    for i, j in zip(sparse_array, [[0, 0], [0, 1]]):
+        assert i == ImmutableSparseNDimArray(j)
+
+    def sparse_assignment():
+        sparse_array[0, 0] = 123
+
+    assert len(sparse_array._sparse_array) == 1
+    raises(TypeError, sparse_assignment)
+    assert len(sparse_array._sparse_array) == 1
+    assert sparse_array[0, 0] == 0
+    assert sparse_array/0 == ImmutableSparseNDimArray([[S.NaN, S.NaN], [S.NaN, S.ComplexInfinity]], (2, 2))
+
+    # test for large scale sparse array
+    # equality test
+    assert ImmutableSparseNDimArray.zeros(100000, 200000) == ImmutableSparseNDimArray.zeros(100000, 200000)
+
+    # __mul__ and __rmul__
+    a = ImmutableSparseNDimArray({200001: 1}, (100000, 200000))
+    assert a * 3 == ImmutableSparseNDimArray({200001: 3}, (100000, 200000))
+    assert 3 * a == ImmutableSparseNDimArray({200001: 3}, (100000, 200000))
+    assert a * 0 == ImmutableSparseNDimArray({}, (100000, 200000))
+    assert 0 * a == ImmutableSparseNDimArray({}, (100000, 200000))
+
+    # __truediv__
+    assert a/3 == ImmutableSparseNDimArray({200001: Rational(1, 3)}, (100000, 200000))
+
+    # __neg__
+    assert -a == ImmutableSparseNDimArray({200001: -1}, (100000, 200000))
+
+
+def test_calculation():
+
+    a = ImmutableDenseNDimArray([1]*9, (3, 3))
+    b = ImmutableDenseNDimArray([9]*9, (3, 3))
+
+    c = a + b
+    for i in c:
+        assert i == ImmutableDenseNDimArray([10, 10, 10])
+
+    assert c == ImmutableDenseNDimArray([10]*9, (3, 3))
+    assert c == ImmutableSparseNDimArray([10]*9, (3, 3))
+
+    c = b - a
+    for i in c:
+        assert i == ImmutableDenseNDimArray([8, 8, 8])
+
+    assert c == ImmutableDenseNDimArray([8]*9, (3, 3))
+    assert c == ImmutableSparseNDimArray([8]*9, (3, 3))
+
+
+def test_ndim_array_converting():
+    dense_array = ImmutableDenseNDimArray([1, 2, 3, 4], (2, 2))
+    alist = dense_array.tolist()
+
+    assert alist == [[1, 2], [3, 4]]
+
+    matrix = dense_array.tomatrix()
+    assert (isinstance(matrix, Matrix))
+
+    for i in range(len(dense_array)):
+        assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
+    assert matrix.shape == dense_array.shape
+
+    assert ImmutableDenseNDimArray(matrix) == dense_array
+    assert ImmutableDenseNDimArray(matrix.as_immutable()) == dense_array
+    assert ImmutableDenseNDimArray(matrix.as_mutable()) == dense_array
+
+    sparse_array = ImmutableSparseNDimArray([1, 2, 3, 4], (2, 2))
+    alist = sparse_array.tolist()
+
+    assert alist == [[1, 2], [3, 4]]
+
+    matrix = sparse_array.tomatrix()
+    assert(isinstance(matrix, SparseMatrix))
+
+    for i in range(len(sparse_array)):
+        assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
+    assert matrix.shape == sparse_array.shape
+
+    assert ImmutableSparseNDimArray(matrix) == sparse_array
+    assert ImmutableSparseNDimArray(matrix.as_immutable()) == sparse_array
+    assert ImmutableSparseNDimArray(matrix.as_mutable()) == sparse_array
+
+
+def test_converting_functions():
+    arr_list = [1, 2, 3, 4]
+    arr_matrix = Matrix(((1, 2), (3, 4)))
+
+    # list
+    arr_ndim_array = ImmutableDenseNDimArray(arr_list, (2, 2))
+    assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray))
+    assert arr_matrix.tolist() == arr_ndim_array.tolist()
+
+    # Matrix
+    arr_ndim_array = ImmutableDenseNDimArray(arr_matrix)
+    assert (isinstance(arr_ndim_array, ImmutableDenseNDimArray))
+    assert arr_matrix.tolist() == arr_ndim_array.tolist()
+    assert arr_matrix.shape == arr_ndim_array.shape
+
+
+def test_equality():
+    first_list = [1, 2, 3, 4]
+    second_list = [1, 2, 3, 4]
+    third_list = [4, 3, 2, 1]
+    assert first_list == second_list
+    assert first_list != third_list
+
+    first_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2))
+    second_ndim_array = ImmutableDenseNDimArray(second_list, (2, 2))
+    fourth_ndim_array = ImmutableDenseNDimArray(first_list, (2, 2))
+
+    assert first_ndim_array == second_ndim_array
+
+    def assignment_attempt(a):
+        a[0, 0] = 0
+
+    raises(TypeError, lambda: assignment_attempt(second_ndim_array))
+    assert first_ndim_array == second_ndim_array
+    assert first_ndim_array == fourth_ndim_array
+
+
+def test_arithmetic():
+    a = ImmutableDenseNDimArray([3 for i in range(9)], (3, 3))
+    b = ImmutableDenseNDimArray([7 for i in range(9)], (3, 3))
+
+    c1 = a + b
+    c2 = b + a
+    assert c1 == c2
+
+    d1 = a - b
+    d2 = b - a
+    assert d1 == d2 * (-1)
+
+    e1 = a * 5
+    e2 = 5 * a
+    e3 = copy(a)
+    e3 *= 5
+    assert e1 == e2 == e3
+
+    f1 = a / 5
+    f2 = copy(a)
+    f2 /= 5
+    assert f1 == f2
+    assert f1[0, 0] == f1[0, 1] == f1[0, 2] == f1[1, 0] == f1[1, 1] == \
+    f1[1, 2] == f1[2, 0] == f1[2, 1] == f1[2, 2] == Rational(3, 5)
+
+    assert type(a) == type(b) == type(c1) == type(c2) == type(d1) == type(d2) \
+        == type(e1) == type(e2) == type(e3) == type(f1)
+
+    z0 = -a
+    assert z0 == ImmutableDenseNDimArray([-3 for i in range(9)], (3, 3))
+
+
+def test_higher_dimenions():
+    m3 = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
+
+    assert m3.tolist() == [[[10, 11, 12, 13],
+            [14, 15, 16, 17],
+            [18, 19, 20, 21]],
+
+           [[22, 23, 24, 25],
+            [26, 27, 28, 29],
+            [30, 31, 32, 33]]]
+
+    assert m3._get_tuple_index(0) == (0, 0, 0)
+    assert m3._get_tuple_index(1) == (0, 0, 1)
+    assert m3._get_tuple_index(4) == (0, 1, 0)
+    assert m3._get_tuple_index(12) == (1, 0, 0)
+
+    assert str(m3) == '[[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]'
+
+    m3_rebuilt = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]])
+    assert m3 == m3_rebuilt
+
+    m3_other = ImmutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]], (2, 3, 4))
+
+    assert m3 == m3_other
+
+
+def test_rebuild_immutable_arrays():
+    sparr = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
+    densarr = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
+
+    assert sparr == sparr.func(*sparr.args)
+    assert densarr == densarr.func(*densarr.args)
+
+
+def test_slices():
+    md = ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
+
+    assert md[:] == ImmutableDenseNDimArray(range(10, 34), (2, 3, 4))
+    assert md[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
+    assert md[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
+    assert md[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
+    assert md[:, :, :] == md
+
+    sd = ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
+    assert sd == ImmutableSparseNDimArray(md)
+
+    assert sd[:] == ImmutableSparseNDimArray(range(10, 34), (2, 3, 4))
+    assert sd[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
+    assert sd[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
+    assert sd[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
+    assert sd[:, :, :] == sd
+
+
+def test_diff_and_applyfunc():
+    from sympy.abc import x, y, z
+    md = ImmutableDenseNDimArray([[x, y], [x*z, x*y*z]])
+    assert md.diff(x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]])
+    assert diff(md, x) == ImmutableDenseNDimArray([[1, 0], [z, y*z]])
+
+    sd = ImmutableSparseNDimArray(md)
+    assert sd == ImmutableSparseNDimArray([x, y, x*z, x*y*z], (2, 2))
+    assert sd.diff(x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]])
+    assert diff(sd, x) == ImmutableSparseNDimArray([[1, 0], [z, y*z]])
+
+    mdn = md.applyfunc(lambda x: x*3)
+    assert mdn == ImmutableDenseNDimArray([[3*x, 3*y], [3*x*z, 3*x*y*z]])
+    assert md != mdn
+
+    sdn = sd.applyfunc(lambda x: x/2)
+    assert sdn == ImmutableSparseNDimArray([[x/2, y/2], [x*z/2, x*y*z/2]])
+    assert sd != sdn
+
+    sdp = sd.applyfunc(lambda x: x+1)
+    assert sdp == ImmutableSparseNDimArray([[x + 1, y + 1], [x*z + 1, x*y*z + 1]])
+    assert sd != sdp
+
+
+def test_op_priority():
+    from sympy.abc import x
+    md = ImmutableDenseNDimArray([1, 2, 3])
+    e1 = (1+x)*md
+    e2 = md*(1+x)
+    assert e1 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x])
+    assert e1 == e2
+
+    sd = ImmutableSparseNDimArray([1, 2, 3])
+    e3 = (1+x)*sd
+    e4 = sd*(1+x)
+    assert e3 == ImmutableDenseNDimArray([1+x, 2+2*x, 3+3*x])
+    assert e3 == e4
+
+
+def test_symbolic_indexing():
+    x, y, z, w = symbols("x y z w")
+    M = ImmutableDenseNDimArray([[x, y], [z, w]])
+    i, j = symbols("i, j")
+    Mij = M[i, j]
+    assert isinstance(Mij, Indexed)
+    Ms = ImmutableSparseNDimArray([[2, 3*x], [4, 5]])
+    msij = Ms[i, j]
+    assert isinstance(msij, Indexed)
+    for oi, oj in [(0, 0), (0, 1), (1, 0), (1, 1)]:
+        assert Mij.subs({i: oi, j: oj}) == M[oi, oj]
+        assert msij.subs({i: oi, j: oj}) == Ms[oi, oj]
+    A = IndexedBase("A", (0, 2))
+    assert A[0, 0].subs(A, M) == x
+    assert A[i, j].subs(A, M) == M[i, j]
+    assert M[i, j].subs(M, A) == A[i, j]
+
+    assert isinstance(M[3 * i - 2, j], Indexed)
+    assert M[3 * i - 2, j].subs({i: 1, j: 0}) == M[1, 0]
+    assert isinstance(M[i, 0], Indexed)
+    assert M[i, 0].subs(i, 0) == M[0, 0]
+    assert M[0, i].subs(i, 1) == M[0, 1]
+
+    assert M[i, j].diff(x) == ImmutableDenseNDimArray([[1, 0], [0, 0]])[i, j]
+    assert Ms[i, j].diff(x) == ImmutableSparseNDimArray([[0, 3], [0, 0]])[i, j]
+
+    Mo = ImmutableDenseNDimArray([1, 2, 3])
+    assert Mo[i].subs(i, 1) == 2
+    Mos = ImmutableSparseNDimArray([1, 2, 3])
+    assert Mos[i].subs(i, 1) == 2
+
+    raises(ValueError, lambda: M[i, 2])
+    raises(ValueError, lambda: M[i, -1])
+    raises(ValueError, lambda: M[2, i])
+    raises(ValueError, lambda: M[-1, i])
+
+    raises(ValueError, lambda: Ms[i, 2])
+    raises(ValueError, lambda: Ms[i, -1])
+    raises(ValueError, lambda: Ms[2, i])
+    raises(ValueError, lambda: Ms[-1, i])
+
+
+def test_issue_12665():
+    # Testing Python 3 hash of immutable arrays:
+    arr = ImmutableDenseNDimArray([1, 2, 3])
+    # This should NOT raise an exception:
+    hash(arr)
+
+
+def test_zeros_without_shape():
+    arr = ImmutableDenseNDimArray.zeros()
+    assert arr == ImmutableDenseNDimArray(0)
+
+def test_issue_21870():
+    a0 = ImmutableDenseNDimArray(0)
+    assert a0.rank() == 0
+    a1 = ImmutableDenseNDimArray(a0)
+    assert a1.rank() == 0

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_mutable_ndim_array.py

@@ -0,0 +1,374 @@
+from copy import copy
+
+from sympy.tensor.array.dense_ndim_array import MutableDenseNDimArray
+from sympy.core.function import diff
+from sympy.core.numbers import Rational
+from sympy.core.singleton import S
+from sympy.core.symbol import Symbol
+from sympy.core.sympify import sympify
+from sympy.matrices import SparseMatrix
+from sympy.matrices import Matrix
+from sympy.tensor.array.sparse_ndim_array import MutableSparseNDimArray
+from sympy.testing.pytest import raises
+
+
+def test_ndim_array_initiation():
+    arr_with_one_element = MutableDenseNDimArray([23])
+    assert len(arr_with_one_element) == 1
+    assert arr_with_one_element[0] == 23
+    assert arr_with_one_element.rank() == 1
+    raises(ValueError, lambda: arr_with_one_element[1])
+
+    arr_with_symbol_element = MutableDenseNDimArray([Symbol('x')])
+    assert len(arr_with_symbol_element) == 1
+    assert arr_with_symbol_element[0] == Symbol('x')
+    assert arr_with_symbol_element.rank() == 1
+
+    number5 = 5
+    vector = MutableDenseNDimArray.zeros(number5)
+    assert len(vector) == number5
+    assert vector.shape == (number5,)
+    assert vector.rank() == 1
+    raises(ValueError, lambda: arr_with_one_element[5])
+
+    vector = MutableSparseNDimArray.zeros(number5)
+    assert len(vector) == number5
+    assert vector.shape == (number5,)
+    assert vector._sparse_array == {}
+    assert vector.rank() == 1
+
+    n_dim_array = MutableDenseNDimArray(range(3**4), (3, 3, 3, 3,))
+    assert len(n_dim_array) == 3 * 3 * 3 * 3
+    assert n_dim_array.shape == (3, 3, 3, 3)
+    assert n_dim_array.rank() == 4
+    raises(ValueError, lambda: n_dim_array[0, 0, 0, 3])
+    raises(ValueError, lambda: n_dim_array[3, 0, 0, 0])
+    raises(ValueError, lambda: n_dim_array[3**4])
+
+    array_shape = (3, 3, 3, 3)
+    sparse_array = MutableSparseNDimArray.zeros(*array_shape)
+    assert len(sparse_array._sparse_array) == 0
+    assert len(sparse_array) == 3 * 3 * 3 * 3
+    assert n_dim_array.shape == array_shape
+    assert n_dim_array.rank() == 4
+
+    one_dim_array = MutableDenseNDimArray([2, 3, 1])
+    assert len(one_dim_array) == 3
+    assert one_dim_array.shape == (3,)
+    assert one_dim_array.rank() == 1
+    assert one_dim_array.tolist() == [2, 3, 1]
+
+    shape = (3, 3)
+    array_with_many_args = MutableSparseNDimArray.zeros(*shape)
+    assert len(array_with_many_args) == 3 * 3
+    assert array_with_many_args.shape == shape
+    assert array_with_many_args[0, 0] == 0
+    assert array_with_many_args.rank() == 2
+
+    shape = (int(3), int(3))
+    array_with_long_shape = MutableSparseNDimArray.zeros(*shape)
+    assert len(array_with_long_shape) == 3 * 3
+    assert array_with_long_shape.shape == shape
+    assert array_with_long_shape[int(0), int(0)] == 0
+    assert array_with_long_shape.rank() == 2
+
+    vector_with_long_shape = MutableDenseNDimArray(range(5), int(5))
+    assert len(vector_with_long_shape) == 5
+    assert vector_with_long_shape.shape == (int(5),)
+    assert vector_with_long_shape.rank() == 1
+    raises(ValueError, lambda: vector_with_long_shape[int(5)])
+
+    from sympy.abc import x
+    for ArrayType in [MutableDenseNDimArray, MutableSparseNDimArray]:
+        rank_zero_array = ArrayType(x)
+        assert len(rank_zero_array) == 1
+        assert rank_zero_array.shape == ()
+        assert rank_zero_array.rank() == 0
+        assert rank_zero_array[()] == x
+        raises(ValueError, lambda: rank_zero_array[0])
+
+def test_sympify():
+    from sympy.abc import x, y, z, t
+    arr = MutableDenseNDimArray([[x, y], [1, z*t]])
+    arr_other = sympify(arr)
+    assert arr_other.shape == (2, 2)
+    assert arr_other == arr
+
+
+def test_reshape():
+    array = MutableDenseNDimArray(range(50), 50)
+    assert array.shape == (50,)
+    assert array.rank() == 1
+
+    array = array.reshape(5, 5, 2)
+    assert array.shape == (5, 5, 2)
+    assert array.rank() == 3
+    assert len(array) == 50
+
+
+def test_iterator():
+    array = MutableDenseNDimArray(range(4), (2, 2))
+    assert array[0] == MutableDenseNDimArray([0, 1])
+    assert array[1] == MutableDenseNDimArray([2, 3])
+
+    array = array.reshape(4)
+    j = 0
+    for i in array:
+        assert i == j
+        j += 1
+
+
+def test_getitem():
+    for ArrayType in [MutableDenseNDimArray, MutableSparseNDimArray]:
+        array = ArrayType(range(24)).reshape(2, 3, 4)
+        assert array.tolist() == [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]]
+        assert array[0] == ArrayType([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
+        assert array[0, 0] == ArrayType([0, 1, 2, 3])
+        value = 0
+        for i in range(2):
+            for j in range(3):
+                for k in range(4):
+                    assert array[i, j, k] == value
+                    value += 1
+
+    raises(ValueError, lambda: array[3, 4, 5])
+    raises(ValueError, lambda: array[3, 4, 5, 6])
+    raises(ValueError, lambda: array[3, 4, 5, 3:4])
+
+
+def test_sparse():
+    sparse_array = MutableSparseNDimArray([0, 0, 0, 1], (2, 2))
+    assert len(sparse_array) == 2 * 2
+    # dictionary where all data is, only non-zero entries are actually stored:
+    assert len(sparse_array._sparse_array) == 1
+
+    assert sparse_array.tolist() == [[0, 0], [0, 1]]
+
+    for i, j in zip(sparse_array, [[0, 0], [0, 1]]):
+        assert i == MutableSparseNDimArray(j)
+
+    sparse_array[0, 0] = 123
+    assert len(sparse_array._sparse_array) == 2
+    assert sparse_array[0, 0] == 123
+    assert sparse_array/0 == MutableSparseNDimArray([[S.ComplexInfinity, S.NaN], [S.NaN, S.ComplexInfinity]], (2, 2))
+
+    # when element in sparse array become zero it will disappear from
+    # dictionary
+    sparse_array[0, 0] = 0
+    assert len(sparse_array._sparse_array) == 1
+    sparse_array[1, 1] = 0
+    assert len(sparse_array._sparse_array) == 0
+    assert sparse_array[0, 0] == 0
+
+    # test for large scale sparse array
+    # equality test
+    a = MutableSparseNDimArray.zeros(100000, 200000)
+    b = MutableSparseNDimArray.zeros(100000, 200000)
+    assert a == b
+    a[1, 1] = 1
+    b[1, 1] = 2
+    assert a != b
+
+    # __mul__ and __rmul__
+    assert a * 3 == MutableSparseNDimArray({200001: 3}, (100000, 200000))
+    assert 3 * a == MutableSparseNDimArray({200001: 3}, (100000, 200000))
+    assert a * 0 == MutableSparseNDimArray({}, (100000, 200000))
+    assert 0 * a == MutableSparseNDimArray({}, (100000, 200000))
+
+    # __truediv__
+    assert a/3 == MutableSparseNDimArray({200001: Rational(1, 3)}, (100000, 200000))
+
+    # __neg__
+    assert -a == MutableSparseNDimArray({200001: -1}, (100000, 200000))
+
+
+def test_calculation():
+
+    a = MutableDenseNDimArray([1]*9, (3, 3))
+    b = MutableDenseNDimArray([9]*9, (3, 3))
+
+    c = a + b
+    for i in c:
+        assert i == MutableDenseNDimArray([10, 10, 10])
+
+    assert c == MutableDenseNDimArray([10]*9, (3, 3))
+    assert c == MutableSparseNDimArray([10]*9, (3, 3))
+
+    c = b - a
+    for i in c:
+        assert i == MutableSparseNDimArray([8, 8, 8])
+
+    assert c == MutableDenseNDimArray([8]*9, (3, 3))
+    assert c == MutableSparseNDimArray([8]*9, (3, 3))
+
+
+def test_ndim_array_converting():
+    dense_array = MutableDenseNDimArray([1, 2, 3, 4], (2, 2))
+    alist = dense_array.tolist()
+
+    assert alist == [[1, 2], [3, 4]]
+
+    matrix = dense_array.tomatrix()
+    assert (isinstance(matrix, Matrix))
+
+    for i in range(len(dense_array)):
+        assert dense_array[dense_array._get_tuple_index(i)] == matrix[i]
+    assert matrix.shape == dense_array.shape
+
+    assert MutableDenseNDimArray(matrix) == dense_array
+    assert MutableDenseNDimArray(matrix.as_immutable()) == dense_array
+    assert MutableDenseNDimArray(matrix.as_mutable()) == dense_array
+
+    sparse_array = MutableSparseNDimArray([1, 2, 3, 4], (2, 2))
+    alist = sparse_array.tolist()
+
+    assert alist == [[1, 2], [3, 4]]
+
+    matrix = sparse_array.tomatrix()
+    assert(isinstance(matrix, SparseMatrix))
+
+    for i in range(len(sparse_array)):
+        assert sparse_array[sparse_array._get_tuple_index(i)] == matrix[i]
+    assert matrix.shape == sparse_array.shape
+
+    assert MutableSparseNDimArray(matrix) == sparse_array
+    assert MutableSparseNDimArray(matrix.as_immutable()) == sparse_array
+    assert MutableSparseNDimArray(matrix.as_mutable()) == sparse_array
+
+
+def test_converting_functions():
+    arr_list = [1, 2, 3, 4]
+    arr_matrix = Matrix(((1, 2), (3, 4)))
+
+    # list
+    arr_ndim_array = MutableDenseNDimArray(arr_list, (2, 2))
+    assert (isinstance(arr_ndim_array, MutableDenseNDimArray))
+    assert arr_matrix.tolist() == arr_ndim_array.tolist()
+
+    # Matrix
+    arr_ndim_array = MutableDenseNDimArray(arr_matrix)
+    assert (isinstance(arr_ndim_array, MutableDenseNDimArray))
+    assert arr_matrix.tolist() == arr_ndim_array.tolist()
+    assert arr_matrix.shape == arr_ndim_array.shape
+
+
+def test_equality():
+    first_list = [1, 2, 3, 4]
+    second_list = [1, 2, 3, 4]
+    third_list = [4, 3, 2, 1]
+    assert first_list == second_list
+    assert first_list != third_list
+
+    first_ndim_array = MutableDenseNDimArray(first_list, (2, 2))
+    second_ndim_array = MutableDenseNDimArray(second_list, (2, 2))
+    third_ndim_array = MutableDenseNDimArray(third_list, (2, 2))
+    fourth_ndim_array = MutableDenseNDimArray(first_list, (2, 2))
+
+    assert first_ndim_array == second_ndim_array
+    second_ndim_array[0, 0] = 0
+    assert first_ndim_array != second_ndim_array
+    assert first_ndim_array != third_ndim_array
+    assert first_ndim_array == fourth_ndim_array
+
+
+def test_arithmetic():
+    a = MutableDenseNDimArray([3 for i in range(9)], (3, 3))
+    b = MutableDenseNDimArray([7 for i in range(9)], (3, 3))
+
+    c1 = a + b
+    c2 = b + a
+    assert c1 == c2
+
+    d1 = a - b
+    d2 = b - a
+    assert d1 == d2 * (-1)
+
+    e1 = a * 5
+    e2 = 5 * a
+    e3 = copy(a)
+    e3 *= 5
+    assert e1 == e2 == e3
+
+    f1 = a / 5
+    f2 = copy(a)
+    f2 /= 5
+    assert f1 == f2
+    assert f1[0, 0] == f1[0, 1] == f1[0, 2] == f1[1, 0] == f1[1, 1] == \
+    f1[1, 2] == f1[2, 0] == f1[2, 1] == f1[2, 2] == Rational(3, 5)
+
+    assert type(a) == type(b) == type(c1) == type(c2) == type(d1) == type(d2) \
+        == type(e1) == type(e2) == type(e3) == type(f1)
+
+    z0 = -a
+    assert z0 == MutableDenseNDimArray([-3 for i in range(9)], (3, 3))
+
+
+def test_higher_dimenions():
+    m3 = MutableDenseNDimArray(range(10, 34), (2, 3, 4))
+
+    assert m3.tolist() == [[[10, 11, 12, 13],
+            [14, 15, 16, 17],
+            [18, 19, 20, 21]],
+
+           [[22, 23, 24, 25],
+            [26, 27, 28, 29],
+            [30, 31, 32, 33]]]
+
+    assert m3._get_tuple_index(0) == (0, 0, 0)
+    assert m3._get_tuple_index(1) == (0, 0, 1)
+    assert m3._get_tuple_index(4) == (0, 1, 0)
+    assert m3._get_tuple_index(12) == (1, 0, 0)
+
+    assert str(m3) == '[[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]]'
+
+    m3_rebuilt = MutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]])
+    assert m3 == m3_rebuilt
+
+    m3_other = MutableDenseNDimArray([[[10, 11, 12, 13], [14, 15, 16, 17], [18, 19, 20, 21]], [[22, 23, 24, 25], [26, 27, 28, 29], [30, 31, 32, 33]]], (2, 3, 4))
+
+    assert m3 == m3_other
+
+
+def test_slices():
+    md = MutableDenseNDimArray(range(10, 34), (2, 3, 4))
+
+    assert md[:] == MutableDenseNDimArray(range(10, 34), (2, 3, 4))
+    assert md[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
+    assert md[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
+    assert md[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
+    assert md[:, :, :] == md
+
+    sd = MutableSparseNDimArray(range(10, 34), (2, 3, 4))
+    assert sd == MutableSparseNDimArray(md)
+
+    assert sd[:] == MutableSparseNDimArray(range(10, 34), (2, 3, 4))
+    assert sd[:, :, 0].tomatrix() == Matrix([[10, 14, 18], [22, 26, 30]])
+    assert sd[0, 1:2, :].tomatrix() == Matrix([[14, 15, 16, 17]])
+    assert sd[0, 1:3, :].tomatrix() == Matrix([[14, 15, 16, 17], [18, 19, 20, 21]])
+    assert sd[:, :, :] == sd
+
+
+def test_slices_assign():
+    a = MutableDenseNDimArray(range(12), shape=(4, 3))
+    b = MutableSparseNDimArray(range(12), shape=(4, 3))
+
+    for i in [a, b]:
+        assert i.tolist() == [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
+        i[0, :] = [2, 2, 2]
+        assert i.tolist() == [[2, 2, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
+        i[0, 1:] = [8, 8]
+        assert i.tolist() == [[2, 8, 8], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
+        i[1:3, 1] = [20, 44]
+        assert i.tolist() == [[2, 8, 8], [3, 20, 5], [6, 44, 8], [9, 10, 11]]
+
+
+def test_diff():
+    from sympy.abc import x, y, z
+    md = MutableDenseNDimArray([[x, y], [x*z, x*y*z]])
+    assert md.diff(x) == MutableDenseNDimArray([[1, 0], [z, y*z]])
+    assert diff(md, x) == MutableDenseNDimArray([[1, 0], [z, y*z]])
+
+    sd = MutableSparseNDimArray(md)
+    assert sd == MutableSparseNDimArray([x, y, x*z, x*y*z], (2, 2))
+    assert sd.diff(x) == MutableSparseNDimArray([[1, 0], [z, y*z]])
+    assert diff(sd, x) == MutableSparseNDimArray([[1, 0], [z, y*z]])

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_ndim_array.py

@@ -0,0 +1,73 @@
+from sympy.testing.pytest import raises
+from sympy.functions.elementary.trigonometric import sin, cos
+from sympy.matrices.dense import Matrix
+from sympy.simplify import simplify
+from sympy.tensor.array import Array
+from sympy.tensor.array.dense_ndim_array import (
+    ImmutableDenseNDimArray, MutableDenseNDimArray)
+from sympy.tensor.array.sparse_ndim_array import (
+    ImmutableSparseNDimArray, MutableSparseNDimArray)
+
+from sympy.abc import x, y
+
+mutable_array_types = [
+    MutableDenseNDimArray,
+    MutableSparseNDimArray
+]
+
+array_types = [
+    ImmutableDenseNDimArray,
+    ImmutableSparseNDimArray,
+    MutableDenseNDimArray,
+    MutableSparseNDimArray
+]
+
+
+def test_array_negative_indices():
+    for ArrayType in array_types:
+        test_array = ArrayType([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
+        assert test_array[:, -1] == Array([5, 10])
+        assert test_array[:, -2] == Array([4, 9])
+        assert test_array[:, -3] == Array([3, 8])
+        assert test_array[:, -4] == Array([2, 7])
+        assert test_array[:, -5] == Array([1, 6])
+        assert test_array[:, 0] == Array([1, 6])
+        assert test_array[:, 1] == Array([2, 7])
+        assert test_array[:, 2] == Array([3, 8])
+        assert test_array[:, 3] == Array([4, 9])
+        assert test_array[:, 4] == Array([5, 10])
+
+        raises(ValueError, lambda: test_array[:, -6])
+        raises(ValueError, lambda: test_array[-3, :])
+
+        assert test_array[-1, -1] == 10
+
+
+def test_issue_18361():
+    A = Array([sin(2 * x) - 2 * sin(x) * cos(x)])
+    B = Array([sin(x)**2 + cos(x)**2, 0])
+    C = Array([(x + x**2)/(x*sin(y)**2 + x*cos(y)**2), 2*sin(x)*cos(x)])
+    assert simplify(A) == Array([0])
+    assert simplify(B) == Array([1, 0])
+    assert simplify(C) == Array([x + 1, sin(2*x)])
+
+
+def test_issue_20222():
+    A = Array([[1, 2], [3, 4]])
+    B = Matrix([[1,2],[3,4]])
+    raises(TypeError, lambda: A - B)
+
+
+def test_issue_17851():
+    for array_type in array_types:
+        A = array_type([])
+        assert isinstance(A, array_type)
+        assert A.shape == (0,)
+        assert list(A) == []
+
+
+def test_issue_and_18715():
+    for array_type in mutable_array_types:
+        A = array_type([0, 1, 2])
+        A[0] += 5
+        assert A[0] == 5

venv/lib/python3.12/site-packages/sympy/tensor/array/tests/test_ndim_array_conversions.py

@@ -0,0 +1,22 @@
+from sympy.tensor.array import (ImmutableDenseNDimArray,
+        ImmutableSparseNDimArray, MutableDenseNDimArray, MutableSparseNDimArray)
+from sympy.abc import x, y, z
+
+
+def test_NDim_array_conv():
+    MD = MutableDenseNDimArray([x, y, z])
+    MS = MutableSparseNDimArray([x, y, z])
+    ID = ImmutableDenseNDimArray([x, y, z])
+    IS = ImmutableSparseNDimArray([x, y, z])
+
+    assert MD.as_immutable() == ID
+    assert MD.as_mutable() == MD
+
+    assert MS.as_immutable() == IS
+    assert MS.as_mutable() == MS
+
+    assert ID.as_immutable() == ID
+    assert ID.as_mutable() == MD
+
+    assert IS.as_immutable() == IS
+    assert IS.as_mutable() == MS