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_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_nested_array_elementwise_add(): cg = _array_contraction( _array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)), (1, 2)) result = _array_add( _array_contraction(_array_tensor_product(M, N), (1, 2)), _array_contraction(_array_tensor_product(N, M), (1, 2))) assert cg == result cg = _array_diagonal( _array_add(_array_tensor_product(M, N), _array_tensor_product(N, M)), (1, 2)) result = _array_add(_array_diagonal(_array_tensor_product(M, N), (1, 2)), _array_diagonal(_array_tensor_product(N, M), (1, 2))) assert cg == result
def _(expr: 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 = [i for i in range(tot1, tot2)] + \ [i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \ [i for i in 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)
def _(expr: ArrayContraction): expr = expr.flatten_contraction_of_diagonal() expr = identify_removable_identity_matrices(expr) expr = expr.split_multiple_contractions() expr = identify_hadamard_products(expr) if not isinstance(expr, ArrayContraction): return _array2matrix(expr) subexpr = expr.expr contraction_indices: tTuple[tTuple[int]] = expr.contraction_indices if contraction_indices == ((0, ), (1, )) or (contraction_indices == ( (0, ), ) and subexpr.shape[1] == 1) or (contraction_indices == ( (1, ), ) and subexpr.shape[0] == 1): shape = subexpr.shape subexpr = _array2matrix(subexpr) if isinstance(subexpr, MatrixExpr): return OneMatrix(1, shape[0]) * subexpr * OneMatrix(shape[1], 1) if isinstance(subexpr, ArrayTensorProduct): newexpr = _array_contraction(_array2matrix(subexpr), *contraction_indices) contraction_indices = newexpr.contraction_indices if any(i > 2 for i in newexpr.subranks): addends = _array_add(*[ _a2m_tensor_product(*j) for j in itertools.product(*[ i.args if isinstance(i, ArrayAdd) else [i] for i in expr.expr.args ]) ]) newexpr = _array_contraction(addends, *contraction_indices) if isinstance(newexpr, ArrayAdd): ret = _array2matrix(newexpr) return ret assert isinstance(newexpr, ArrayContraction) ret = _support_function_tp1_recognize(contraction_indices, list(newexpr.expr.args)) return ret elif not isinstance(subexpr, _CodegenArrayAbstract): ret = _array2matrix(subexpr) if isinstance(ret, MatrixExpr): assert expr.contraction_indices == ((0, 1), ) return _a2m_trace(ret) else: return _array_contraction(ret, *expr.contraction_indices)
def test_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_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 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.conv_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(set( [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, ()
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
def _(expr: ArrayAdd, x: Expr): return _array_add(*[array_derive(arg, x) for arg in expr.args])
def _a2m_add(*args): if not any(isinstance(i, _CodegenArrayAbstract) for i in args): from sympy.matrices.expressions.matadd import MatAdd return MatAdd(*args).doit() else: return _array_add(*args)