From dfab760b4a328d9fa29cef123e0fe8e2926b0c8c Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Fri, 12 Oct 2018 20:57:00 -0700 Subject: ENH: __array_function__ for np.einsum and np.block --- numpy/core/shape_base.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) (limited to 'numpy/core/shape_base.py') diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py index fde23076b..69e8c0afc 100644 --- a/numpy/core/shape_base.py +++ b/numpy/core/shape_base.py @@ -499,7 +499,16 @@ def _block(arrays, max_depth, result_ndim, depth=0): return _atleast_nd(arrays, result_ndim) -# TODO: support array_function_dispatch +def _block_dispatcher(arrays): + 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. -- cgit v1.2.1