Merge pull request #22668 from WarrenWeckesser/use-keepdims

DOC: lib: Use keepdims in a couple docstrings.

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.