Repository: sympy/sympy
File: sympy/tensor/array/expressions/arrayexpr_derivatives.py (Python)

Code before:
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):
    if expr.free_symbols & x.free_symbols:
        raise NotImplementedError("algorithm not implemented for this case")
    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:

Review comment (suggestion):
Here NotImplementedError is not caught by anything. Should this not somehow give an unevaluated derivative rather than raising?

```python
In [2]: D = symbols('D')
   ...: w = S.MatrixSymbol("w", D, 1)

In [3]: w[0,0].diff(w)
---------------------------------------------------------------------------
NotImplementedError
```
https://github.com/sympy/sympy/issues/24725#issuecomment-1468099499

Code after:
from sympy.matrices.expressions.special import Identity, OneMatrix, MatrixUnit
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):
    if expr.free_symbols & x.free_symbols:
        if isinstance(expr, MatrixElement) and isinstance(x, MatrixSymbol):
            return MatrixUnit(x.shape[0], x.shape[1], expr.i, expr.j)
        raise NotImplementedError("algorithm not implemented for this case")
    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)