How to sum an Array or ArraySymbol over an axis?

[redeboer]

Moving the ideas mentioned by @ThePauliPrinciple in #22265 (comment) to this separate discussion. See also the indexing ideas mentioned by @ThePauliPrinciple in #22219.

It seems that it is not yet possible to perform computations on the axes of (symbolic) arrays. This would be required for creating expressions that can be lambdified to numpy.sum(..., axis=1).

To simplify things, let's stick with this sum as an example. We would like to do something like this:

>>> from sympy.tensor.array.expressions.array_expressions import ArraySymbol
>>> A = ArraySymbol("A")
>>> expr = A.sum(axis=1)  # or some other syntax; probably requires a new Sum class
>>>
>>> from sympy.printing.numpy import NumPyPrinter
>>> NumPyPrinter().doprint(expr)
'numpy.sum(A, axis=1)'
>>>
>>> from sympy.printing.tensorflow import TensorflowPrinter
>>> TensorflowPrinter().doprint(expr)
'tf.math.reduce_sum(A, axis=1)'

Is this perhaps already possible? Or are there ideas for some implementation?

[ThePauliPrinciple]

Instead of a new Sum class A.sum(axis=1) could return an ArraySymbol which "remembers" that it still needs to sum over some index (but no longer allows accessing it), then the printer for ArraySymbol can take care of of printing the sum statement. I do like to point out though, that there would not be an easy translation for e.g. c code, where such a sum could require an assignment to a temporary variable.

[ThePauliPrinciple]

You are correct that this should be added to the TensorflowPrinter instead before merging, but this allows for testing new objects without integrating them into the library.

Would you consider ArraySlice to allow None (both numpy and tensorflow allow this as an alternative to newaxis), ArraySymbols and Arrays as valid slices? to replicate this behaviour:

[redeboer]

What's wrong with the existing Sum class?

Do you mean for instance this?

import sympy as sp
A = sp.ArraySymbol("A")
i = sp.Symbol("i")
sp.Sum(A[:, i], (i, 0, sp.oo))

Fine from a symbolic point of view, but can't lambdify to a vectorised form like numpy.sum. (In that sense, the discussion is perhaps too much oriented to code generation/lambdifying.)

[redeboer]

#22269 (reply in thread)
@ThePauliPrinciple thanks for the question, could you (re)post it in #22265?

[oscarbenjamin]

See also #5642 (comment)

[ThePauliPrinciple]

I feel that sum could work:

from sympy import Sum,Symbol,oo
i=Symbol('i')

class modNumPyPrinter(NumPyPrinter):
  def _print_Sum(printer,expr):
    print(expr)
    return ''

modNumPyPrinter().doprint(Sum(A[:,i,:],(i,0,oo)))

_print_Sum could be written to detect if there are any ArraySlices with the symbol over which it sums, then extract which axis this corresponds to.

This does require a symbol to be a valid axis in ArraySlice (similar to an integer)

[redeboer]

I think einsum can be correctly used for summing over the axis of a single Array. E.g.

import numpy as np
A=np.ones((3,))
np.einsum('i->',A)

Would yield a scalar=3. I would imagine that it would shortcut to a special case and just run np.sum if it detects this is possible.

One could also imagine something like i=Array.get_axis(1) and then Sum(Array,i) could lead to printing np.sum(Array,axis=1).

If you need any help with the printing routines, or want me to give them an attempt that is no problem. I'm not very experienced with the Array part of sympy though, so I'm looking for your help on that side e.g giving an example of what sympy expression should lead to what print statement, or even simply the sympy expression (if it is clear what it should mean) that should be supported for printing.
#22265 (comment)

I'm not too familiar with einsum, but it seems one can use it to sum over specific axes as well:

>>> import numpy as np
>>> A = np.arange(9).reshape(3, 3); A
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>>
>>> np.einsum("ij->i", A)  # np.sum(A, axis=1)
array([ 3, 12, 21])
>>> np.einsum("ji->i", A)  # np.sum(A, axis=0)
array([ 9, 12, 15])
>>>
>>> np.einsum("ij->i", A.transpose())
array([ 9, 12, 15])
>>> np.einsum(A, [0, 1], [0])  # np.sum(A, axis=1)
array([ 3, 12, 21])

The last one was added, because perhaps something can be done with a combination of PermuteDims (which lambdifies to transpose) and ArrayContraction (which lambdifies to einsum). Note that ArrayContraction lambdifies einsum with the index notation as shown in the end.

sympy/sympy/printing/numpy.py

Lines 270 to 274 in 3abd23d

	elems = ["%s, %s" % (self._print(arg), ind) for arg, ind in zip(base.args, indices)]
	return "%s(%s)" % (
	self._module_format(self._module + '.einsum'),
	", ".join(elems)
	)

[redeboer]

Ah ok thanks a lot. Yeah I'm not too familiar with tensor/array computations, so glad to read your thoughts on this.

At any rate, it seems clear that some class like ArraySum would have to be defined to get this to work.

[ThePauliPrinciple]

So currently, with my "Indexed story", I can obtain what you are looking for in the example:

But here it is the assignment class which takes care of the book keeping for the sum.

[Upabjojr]

I think the problem with ArrayContraction is that it expects a pair of indices to be contracted on (which is the only sensible contraction for a Tensor, but not neccesarily for an Array and precludes summing over a single index)

We are not limited to contractions over two axes. That makes sense in differential geometry because you want to create invariants independent of the choice of the coordinate system.

However, there are other use cases. For example, in statistics, if you compute the gradient of a neural network, you may get a triple-index contraction.

[Upabjojr]

Note that with lambdify, it is expected that it can work with an arbitrary number of additional dimensions.
Fortunately, this can be done with einsum since it accepts an ellipsis e.g. "...ij->...i".

IIRC, in NumPy (but not in TensorFlow), einsum also accepts tuples of axis integer-valued indices.

[Upabjojr]

Another useful reference (may or may not be helpful): https://optimized-einsum.readthedocs.io/en/stable/#

[Upabjojr]

It seems that it is not yet possible to perform computations on the axes of (symbolic) arrays. This would be required for creating expressions that can be lambdified to numpy.sum(..., axis=1).

ArrayContraction can be given single axes:

ArrayContraction(..., (1,))

This is equivalent to numpy.sum(..., axis=1)

[Upabjojr]

Currently there is an auto-flattening mode (part of the normalization):

In [3]: A = ArraySymbol("A", 3, 3, 3, 3)

In [4]: ArrayContraction(ArrayContraction(A, (2,)), (0, 2))
Out[4]: ArrayContraction(A, (0, 3), (2,))

Maybe this should be disabled. For the code printer, it would be more convenient to separate multiple contractions from single contractions, and then handle the two nest ArrayContraction objects with different code printing logics. That is, reversing the normalization operation.

[ThePauliPrinciple]

I would argue the opposite, the code printer could always recursively call itself whenever multiple contractions are detected and only write one contraction (popping one contraction every time) at a time if the language associated with that printer does not support summing over multiple indices in one call.

[Upabjojr]

Well, the idea is that you should be able to restructure the array expressions inside SymPy before passing it to the code printer. Currently this is not possible because ArrayContraction had automatic normalization by default and that cannot be disabled right now. So you will always have a standard contraction form.

It would make sense to apply some transformation rules before code printing, in order to create an optimized expression tree inside SymPy first, and then call the printer.

[ThePauliPrinciple]

Sorry, my point is only about the default behaviour, not to not have the option at all.

Would differently nested array contractions still lead to a mathematical identity in that case? (or zero when subtracted) As far as I understand, the order of contraction does not matter (although the index numbers might be different if the order is different).

Of course, I understand that the order of contraction might matter (several orders of magnitude) when computing the contraction if it involves an array tensor product for the computational effort.

Creating the optimized expression tree might be more appropriate in codegen.ast. The option to not have normalization would certainly be useful, but it feels to me that SymPy's default behaviour should be most convenient for the mathematical aspect and not the codegen aspect.

Reusing einsumpath from numpy might be helpful (requiring the user to give some approximation of the array sizes).

[Upabjojr]

As far as I understand, the order of contraction does not matter (although the index numbers might be different if the order is different).

In a single ArrayContraction object, that is true and the index order does not matter. In case of nested contractions, the outer ArrayContraction object has less axes, so the contraction indices have to be renumbered (and trust me, renumbering the axes is not as simple as it sounds, though the algorithm is already implemented by now).

The point is, you may wish to have a NumPy code printer that either prints np.einsum or np.sum or both depending on whether the contractions are on multiple axes or on a single axis. The contraction axes have to be renumbered in either np.einsum or np.sum, depending on which on is the outer one.

Creating the optimized expression tree might be more appropriate in codegen.ast.

Well, sympy.tensor.array.expressions is meant to be an AST for array expressions. More generally, all of SymPy is an AST for mathematical concepts.

The option to not have normalization would certainly be useful, but it feels to me that SymPy's default behaviour should be most convenient for the mathematical aspect and not the codegen aspect.

I added default normalization to array expressions for simplicity reasons during the development. It should not be the default. There are three functions (tensorcontraction, permutedims, tensordiagonal) that should be used instead of calling the corresponding class constructors (ArrayContraction, PermuteDims, ArrayDiagonal). I haven't tested them extensively, so I'm not sure they always work, but my idea is that the functions should perform normalizations, while the class constructors shouldn't.

[Upabjojr]

For reference, I decided to base the class ArrayContraction on a similar idea from Wolfram Mathematica. Mathematica's API seems to me clearer and easier than NumPy's one in some cases.