Repository: dask/dask

Issue:
da.einsum ignores split_every
As of git head (495a3611c1ccc12ccc37cf8a56ec3a88743815f5), da.einsum accepts, but ignores, the split_every parameter:
```
import dask.array as da
a = da.ones((5, 40), chunks=10)
b = da.ones((40, ), chunks=10)
da.einsum('...i,...i', a, b, split_every=2).visualize()
```
![download](https://user-images.githubusercontent.com/6213168/39653806-5310b6f8-4fea-11e8-9ba2-a0379299194b.png)


Discussion:
crusaderky: ```
import dask.array as da
a = da.ones((5, 40), chunks=10)
b = da.ones((40, ), chunks=10)
da.einsum('...i,...i', a, b, split_every=2).visualize()
```
![download 1](https://user-images.githubusercontent.com/6213168/39655885-29d298d8-4ff4-11e8-80db-b112c4b8894d.png)

crusaderky: Ready for final review and merge
crusaderky: @shoyer @mrocklin apologies for the delays. Explicit kwargs done as requested.

Patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
index 2879f38e5..fe4d64f76 100644
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -235,3 +235,14 @@ def view(x, dtype, order='C'):
     else:
         x = np.asfortranarray(x)
         return x.T.view(dtype).T
+
+
+def einsum(*operands, **kwargs):
+    subscripts = kwargs.pop('subscripts')
+    ncontract_inds = kwargs.pop('ncontract_inds')
+    dtype = kwargs.pop('kernel_dtype')
+    chunk = np.einsum(subscripts, *operands, dtype=dtype, **kwargs)
+
+    # Avoid concatenate=True in atop by adding 1's
+    # for the contracted dimensions
+    return chunk.reshape(chunk.shape + (1,) * ncontract_inds)
diff --git a/dask/array/einsumfuncs.py b/dask/array/einsumfuncs.py
index 483335a1d..f32a5b7c7 100644
--- a/dask/array/einsumfuncs.py
+++ b/dask/array/einsumfuncs.py
@@ -7,6 +7,7 @@ import numpy as np
 from numpy.compat import basestring
 
 from .core import (atop, asarray)
+from . import chunk
 
 einsum_symbols = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 einsum_symbols_set = set(einsum_symbols)
@@ -182,24 +183,20 @@ def parse_einsum_input(operands):
     return (input_subscripts, output_subscript, operands)
 
 
-def _einsum_kernel(*operands, **kwargs):
-    subscripts = kwargs.pop('subscripts')
-    ncontract_inds = kwargs.pop('ncontract_inds')
-    dtype = kwargs.pop('kernel_dtype')
-    chunk = np.einsum(subscripts, *operands, dtype=dtype, **kwargs)
-
-    # Avoid concatenate=True in atop by adding 1's
-    # for the contracted dimensions
-    return chunk.reshape(chunk.shape + (1,) * ncontract_inds)
-
-
 einsum_can_optimize = LooseVersion(np.__version__) >= LooseVersion("1.12.0")
 
 
 @wraps(np.einsum)
 def einsum(*operands, **kwargs):
-    dtype = kwargs.get('dtype')
-    optimize = kwargs.get('optimize')
+    casting = kwargs.pop('casting', 'safe')
+    dtype = kwargs.pop('dtype', None)
+    optimize = kwargs.pop('optimize', False)
+    order = kwargs.pop('order', 'K')
+    split_every = kwargs.pop('split_every', None)
+    if kwargs:
+        raise TypeError("einsum() got unexpected keyword "
+                        "argument(s) %s" % ",".join(kwargs))
+
     einsum_dtype = dtype
 
     inputs, outputs, ops = parse_einsum_input(operands)
@@ -209,16 +206,18 @@ def einsum(*operands, **kwargs):
     if dtype is None:
         dtype = np.result_type(*[o.dtype for o in ops])
 
-    if optimize is None:
-        optimize = False
-
-    if einsum_can_optimize and optimize is not False:
-        # Avoid computation of dask arrays within np.einsum_path
-        # by passing in small numpy arrays broadcasted
-        # up to the right shape
-        fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape)
-                    for o in ops]
-        optimize, _ = np.einsum_path(subscripts, *fake_ops, optimize=optimize)
+    if einsum_can_optimize:
+        if optimize is not False:
+            # Avoid computation of dask arrays within np.einsum_path
+            # by passing in small numpy arrays broadcasted
+            # up to the right shape
+            fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape)
+                        for o in ops]
+            optimize, _ = np.einsum_path(subscripts, *fake_ops,
+                                         optimize=optimize)
+        kwargs = {'optimize': optimize}
+    else:
+        kwargs = {}
 
     inputs = [tuple(i) for i in inputs.split(",")]
 
@@ -229,27 +228,21 @@ def einsum(*operands, **kwargs):
     contract_inds = all_inds - set(outputs)
     ncontract_inds = len(contract_inds)
 
-    # Update kwargs with np.einsum parameters
-    kwargs['subscripts'] = subscripts
-    kwargs['kernel_dtype'] = einsum_dtype
-    kwargs['ncontract_inds'] = ncontract_inds
-
-    if einsum_can_optimize:
-        kwargs['optimize'] = optimize
-
-    # Update kwargs with atop parameters
-    kwargs['adjust_chunks'] = {ind: 1 for ind in contract_inds}
-    kwargs['dtype'] = dtype
-
     # Introduce the contracted indices into the atop product
     # so that we get numpy arrays, not lists
-    result = atop(_einsum_kernel, tuple(outputs) + tuple(contract_inds),
+    result = atop(chunk.einsum, tuple(outputs) + tuple(contract_inds),
                   *(a for ap in zip(ops, inputs) for a in ap),
-                  **kwargs)
+                  # atop parameters
+                  adjust_chunks={ind: 1 for ind in contract_inds}, dtype=dtype,
+                  # np.einsum parameters
+                  subscripts=subscripts, kernel_dtype=einsum_dtype,
+                  ncontract_inds=ncontract_inds, order=order,
+                  casting=casting, **kwargs)
 
     # Now reduce over any extra contraction dimensions
     if ncontract_inds > 0:
         size = len(outputs)
-        return result.sum(axis=list(range(size, size + ncontract_inds)))
+        return result.sum(axis=list(range(size, size + ncontract_inds)),
+                          split_every=split_every)
 
     return result
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index ae9245a15..1d1657aa7 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -9,6 +9,7 @@ Array
 +++++
 
 - Fix ``rechunk`` with chunksize of -1 in a dict (:pr:`3469`) `Stephan Hoyer`_
+- ``einsum`` now accepts the ``split_every`` parameter (:pr:`3396`) `Guido Imperiale`_
 
 Dataframe
 +++++++++
diff --git a/docs/source/delayed.rst b/docs/source/delayed.rst
index 4cfd86f86..63bff2b9b 100644
--- a/docs/source/delayed.rst
+++ b/docs/source/delayed.rst
@@ -69,7 +69,7 @@ execution, placing the function and its arguments into a task graph.
 .. autosummary::
     delayed
 
-We slightly modify our code our code by wrapping functions in ``delayed``.
+We slightly modify our code by wrapping functions in ``delayed``.
 This delays the execution of the function and generates a dask graph instead.
 
 .. code-block:: python


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index a5a214eab..f90291102 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1407,6 +1407,19 @@ def test_einsum_casting(casting):
               da.einsum(sig, *np_inputs, casting=casting))
 
 
+@pytest.mark.parametrize('split_every', [None, 2])
+def test_einsum_split_every(split_every):
+    np_inputs, da_inputs = _numpy_and_dask_inputs('a')
+    assert_eq(np.einsum('a', *np_inputs),
+              da.einsum('a', *da_inputs, split_every=split_every))
+
+
+def test_einsum_invalid_args():
+    _, da_inputs = _numpy_and_dask_inputs('a')
+    with pytest.raises(TypeError):
+        da.einsum('a', *da_inputs, foo=1, bar=2)
+
+
 def test_einsum_broadcasting_contraction():
     a = np.random.rand(1, 5, 4)
     b = np.random.rand(4, 6)