diff options
Diffstat (limited to 'numpy/lib/shape_base.py')
| -rw-r--r-- | numpy/lib/shape_base.py | 24 |
1 files changed, 13 insertions, 11 deletions
diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py index ab91423d9..3cb4dc19c 100644 --- a/numpy/lib/shape_base.py +++ b/numpy/lib/shape_base.py @@ -1,9 +1,7 @@ import functools import numpy.core.numeric as _nx -from numpy.core.numeric import ( - asarray, zeros, outer, concatenate, array, asanyarray - ) +from numpy.core.numeric import asarray, zeros, array, asanyarray from numpy.core.fromnumeric import reshape, transpose from numpy.core.multiarray import normalize_axis_index from numpy.core import overrides @@ -124,19 +122,21 @@ def take_along_axis(arr, indices, axis): >>> np.sort(a, axis=1) array([[10, 20, 30], [40, 50, 60]]) - >>> ai = np.argsort(a, axis=1); ai + >>> ai = np.argsort(a, axis=1) + >>> ai array([[0, 2, 1], [1, 2, 0]]) >>> np.take_along_axis(a, ai, axis=1) array([[10, 20, 30], [40, 50, 60]]) - The same works for max and min, if you expand the dimensions: + The same works for max and min, if you maintain the trivial dimension + with ``keepdims``: - >>> np.expand_dims(np.max(a, axis=1), axis=1) + >>> np.max(a, axis=1, keepdims=True) array([[30], [60]]) - >>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1) + >>> ai = np.argmax(a, axis=1, keepdims=True) >>> ai array([[1], [0]]) @@ -147,8 +147,8 @@ def take_along_axis(arr, indices, axis): If we want to get the max and min at the same time, we can stack the indices first - >>> ai_min = np.expand_dims(np.argmin(a, axis=1), axis=1) - >>> ai_max = np.expand_dims(np.argmax(a, axis=1), axis=1) + >>> ai_min = np.argmin(a, axis=1, keepdims=True) + >>> ai_max = np.argmax(a, axis=1, keepdims=True) >>> ai = np.concatenate([ai_min, ai_max], axis=1) >>> ai array([[0, 1], @@ -237,7 +237,7 @@ def put_along_axis(arr, indices, values, axis): We can replace the maximum values with: - >>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1) + >>> ai = np.argmax(a, axis=1, keepdims=True) >>> ai array([[1], [0]]) @@ -372,7 +372,7 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs): # invoke the function on the first item try: ind0 = next(inds) - except StopIteration as e: + except StopIteration: raise ValueError( 'Cannot apply_along_axis when any iteration dimensions are 0' ) from None @@ -1041,6 +1041,7 @@ def dsplit(ary, indices_or_sections): raise ValueError('dsplit only works on arrays of 3 or more dimensions') return split(ary, indices_or_sections, 2) + def get_array_prepare(*args): """Find the wrapper for the array with the highest priority. @@ -1053,6 +1054,7 @@ def get_array_prepare(*args): return wrappers[-1][-1] return None + def get_array_wrap(*args): """Find the wrapper for the array with the highest priority. |