Repository: chainer/chainer

Issue:
Add `F.einsum`
See also #4016.

This PR implements einsum as a single `FunctionNode`. This approach is different from https://github.com/chainer/chainer/issues/4016#issue-278690083.

If three or more operands are given, as discussed `Prod.backward` somewhere, I believe that the backward would be faster if we separated FunctionNode into steps of `numpy.einsum_path` (c.f. `optimize` option of np.einsum).
It's possible that sum operations in a single operand (e.g. `b` in `einsum('ab,ac->c', x, y)`) should be preprocessed alone for the same reason.


Discussion:
We should be aware that there is an inconsistency between syntax of `einsum` subscripts between C parser and Python parser (which is contained in this PR).
https://github.com/numpy/numpy/issues/10926
On `einsum` subscripts, the permissive syntax is natural to me and the parser in this PR is based on the "newer" parser.  (I further generalized it to a kind of "diag einsum" syntax.)
The issue numpy/numpy#10926 was avoided at the commit 6e176507b6dbb14ab7b959c3bfb383e1757eb406 in this PR, probably.
With the patch cupy/cupy#1181, `python -m pytest chainer_tests/functions_tests/math_tests/test_einsum.py` passes (including gpu tests).
The PR requires fast `cupy.einsum`.
@toslunar Could you tell me what PR blocks this PR?
- cupy/cupy#1218 blocks this PR.  This `F.einsum` shouldn't be merged until `cupy.einsum` becomes reasonably fast.
- This PR needs work.  `numpy.einsum` (with default `optimize=False` before the next version 1.15) with 3 or more operands is also slow.  To support faster einsum in CPU, this PR should be totally rewritten as I commented:
  > If three or more operands are given, as discussed Prod.backward somewhere, I believe that the backward would be faster if we separated FunctionNode into steps of numpy.einsum_path (c.f. optimize option of np.einsum).
It's possible that sum operations in a single operand (e.g. b in einsum('ab,ac->c', x, y)) should be preprocessed alone for the same reason.
https://github.com/chainer/chainer/pull/4644#issue-182088644
Thank you for clarification!

> This PR needs work. numpy.einsum (with default optimize=False before the next version 1.15) with 3 or more operands is also slow. To support faster einsum in CPU, this PR should be totally rewritten as I commented:

Do you mean you don't need this work if `optimize=True` can be used? I'm not sure how much is that work, but I think it is also acceptable to just use `optimize=True` for NumPy 1.12+ and show warning when older NumPy (that does not support optimize) is used.
I agree with that.  Showing warning message in numpy<1.12 is fine to me.  The further improvement can be in a new PR, possibly with numpy-compatible broadcast feature.
cupy/cupy#1218 is merged
jenkins, test this please.
Jenkins CI test (for commit ccc64c9cfdf280caa01ad7a78d695448f44dd1fd) succeeded without errors!

Jenkins, test this please.
Jenkins CI test (for commit b33d4b5c40499b0b84ffcc34dfb5816c2c7cfc14) failed with status FAILURE.
(For contributors, please wait until the reviewer confirms the details of the error.)

Jenkins, test this please.
Jenkins CI test (for commit cdf72d7d6fdd21c2604e26fd3d5a6769519982f7) succeeded without errors!

The Travis failure (`TestL2Normalization_param_18.test_backward_cpu`) is irrelevant to the PR.
LGTM. Please resolve conflicts.
jenkins, test this please.
Jenkins CI test (for commit 56d15a84c2bc748db0fc0914ae4244baa9e4062c) succeeded without errors!

LGTM!

Resolving patch:
diff --git a/chainer/functions/math/einsum.py b/chainer/functions/math/einsum.py
--- a/chainer/functions/math/einsum.py
+++ b/chainer/functions/math/einsum.py
@@ -164,7 +164,7 @@ def forward(self, inputs):
                 diag_y.setflags(write=True)
             diag_y[...] = _einsum(
                 xp, dtype, self.in_subs, ''.join(direct_sub), *inputs
-            )[expander]
+            )[tuple(expander)]
         return y,
 
 


Test patch: