numpy · mhvk · Oct 26, 2018 · Oct 13, 2018 · Oct 19, 2018 · Oct 25, 2018
diff --git a/numpy/core/einsumfunc.py b/numpy/core/einsumfunc.py
@@ -9,6 +9,7 @@
 from numpy.compat import basestring
 from numpy.core.multiarray import c_einsum
 from numpy.core.numeric import asanyarray, tensordot
+from numpy.core.overrides import array_function_dispatch
 
 __all__ = ['einsum', 'einsum_path']
 
@@ -689,6 +690,17 @@ def _parse_einsum_input(operands):
     return (input_subscripts, output_subscript, operands)
 
 
+def _einsum_path_dispatcher(*operands, **kwargs):
+    # NOTE: technically, we should only dispatch on array-like arguments, not
+    # subscripts (given as strings). But separating operands into
+    # arrays/subscripts is a little tricky/slow (given einsum's two supported
+    # signatures), so as a practical shortcut we dispatch on everything.
+    # Strings will be ignored for dispatching since they don't define
+    # __array_function__.
+    return operands
+
+
+@array_function_dispatch(_einsum_path_dispatcher)
 def einsum_path(*operands, **kwargs):
     """
     einsum_path(subscripts, *operands, optimize='greedy')
@@ -980,7 +992,16 @@ def einsum_path(*operands, **kwargs):
     return (path, path_print)
 
 
+def _einsum_dispatcher(*operands, **kwargs):
+    # Arguably we dispatch on more arguments that we really should; see note in
+    # _einsum_path_dispatcher for why.
+    for op in operands:
+        yield op
+    yield kwargs.get('out')
+
+
 # Rewrite einsum to handle different cases
+@array_function_dispatch(_einsum_dispatcher)
 def einsum(*operands, **kwargs):
     """
     einsum(subscripts, *operands, out=None, dtype=None, order='K',

diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
@@ -631,7 +631,19 @@ def _block(arrays, max_depth, result_ndim, depth=0):
         return _atleast_nd(arrays, result_ndim)
 
 
-# TODO: support array_function_dispatch
+def _block_dispatcher(arrays):
+    # Use type(...) is list to match the behavior of np.block(), which special
+    # cases list specifically rather than allowing for generic iterables or
+    # tuple. Also, we know that list.__array_function__ will never exist.
+    if type(arrays) is list:
+        for subarrays in arrays:
+            for subarray in _block_dispatcher(subarrays):
+                yield subarray
+    else:
+        yield arrays
+
+
+@array_function_dispatch(_block_dispatcher)
 def block(arrays):
     """
     Assemble an nd-array from nested lists of blocks.

diff --git a/numpy/core/tests/test_shape_base.py b/numpy/core/tests/test_shape_base.py
@@ -6,8 +6,7 @@
     array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack,
     newaxis, concatenate, stack
     )
-
-from numpy.core.shape_base import (_block_setup,
+from numpy.core.shape_base import (_block_dispatcher, _block_setup,
                                    _block_concatenate, _block_slicing)
 from numpy.testing import (
     assert_, assert_raises, assert_array_equal, assert_equal,
@@ -677,3 +676,17 @@ def test_block_memory_order(self, block):
 
         assert block(b_c).flags['C_CONTIGUOUS']
         assert block(b_f).flags['F_CONTIGUOUS']
+
+
+def test_block_dispatcher():
+    class ArrayLike(object):
+        pass
+    a = ArrayLike()
+    b = ArrayLike()
+    c = ArrayLike()
+    assert_equal(list(_block_dispatcher(a)), [a])
+    assert_equal(list(_block_dispatcher([a])), [a])
+    assert_equal(list(_block_dispatcher([a, b])), [a, b])
+    assert_equal(list(_block_dispatcher([[a], [b, [c]]])), [a, b, c])
+    # don't recurse into non-lists
+    assert_equal(list(_block_dispatcher((a, b))), [(a, b)])