diff options
Diffstat (limited to 'numpy/core/shape_base.py')
| -rw-r--r-- | numpy/core/shape_base.py | 22 |
1 files changed, 15 insertions, 7 deletions
diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py index c23ffa935..710f64827 100644 --- a/numpy/core/shape_base.py +++ b/numpy/core/shape_base.py @@ -123,7 +123,7 @@ def atleast_2d(*arys): if ary.ndim == 0: result = ary.reshape(1, 1) elif ary.ndim == 1: - result = ary[newaxis,:] + result = ary[newaxis, :] else: result = ary res.append(result) @@ -193,9 +193,9 @@ def atleast_3d(*arys): if ary.ndim == 0: result = ary.reshape(1, 1, 1) elif ary.ndim == 1: - result = ary[newaxis,:, newaxis] + result = ary[newaxis, :, newaxis] elif ary.ndim == 2: - result = ary[:,:, newaxis] + result = ary[:, :, newaxis] else: result = ary res.append(result) @@ -432,6 +432,14 @@ def stack(arrays, axis=0, out=None): return _nx.concatenate(expanded_arrays, axis=axis, out=out) +# Internal functions to eliminate the overhead of repeated dispatch in one of +# the two possible paths inside np.block. +# Use getattr to protect against __array_function__ being disabled. +_size = getattr(_nx.size, '__wrapped__', _nx.size) +_ndim = getattr(_nx.ndim, '__wrapped__', _nx.ndim) +_concatenate = getattr(_nx.concatenate, '__wrapped__', _nx.concatenate) + + def _block_format_index(index): """ Convert a list of indices ``[0, 1, 2]`` into ``"arrays[0][1][2]"``. @@ -512,8 +520,8 @@ def _block_check_depths_match(arrays, parent_index=[]): return parent_index + [None], 0, 0 else: # We've 'bottomed out' - arrays is either a scalar or an array - size = _nx.size(arrays) - return parent_index, _nx.ndim(arrays), size + size = _size(arrays) + return parent_index, _ndim(arrays), size def _atleast_nd(a, ndim): @@ -656,7 +664,7 @@ def _block(arrays, max_depth, result_ndim, depth=0): if depth < max_depth: arrs = [_block(arr, max_depth, result_ndim, depth+1) for arr in arrays] - return _nx.concatenate(arrs, axis=-(max_depth-depth)) + return _concatenate(arrs, axis=-(max_depth-depth)) else: # We've 'bottomed out' - arrays is either a scalar or an array # type(arrays) is not list @@ -874,7 +882,7 @@ def _block_slicing(arrays, list_ndim, result_ndim): # Test preferring F only in the case that all input arrays are F F_order = all(arr.flags['F_CONTIGUOUS'] for arr in arrays) - C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays) + C_order = all(arr.flags['C_CONTIGUOUS'] for arr in arrays) order = 'F' if F_order and not C_order else 'C' result = _nx.empty(shape=shape, dtype=dtype, order=order) # Note: In a c implementation, the function |