MAINT: Add default value for axis parameter in numpy.take_along_axis

mtsokol · mtsokol · commit ff92ba826eb9 · 2025-05-10T09:27:18.000Z
diff --git a/doc/release/upcoming_changes/28615.change.rst b/doc/release/upcoming_changes/28615.change.rst
@@ -1,3 +1,5 @@
 * NumPy's ``__array_api_version__`` was upgraded from ``2023.12`` to ``2024.12``.
 * `numpy.count_nonzero` for ``axis=None`` (default) now returns a NumPy scalar
   instead of a Python integer.
+* The parameter ``axis`` in `numpy.take_along_axis` function has now a default
+  value of ``-1``.
diff --git a/numpy/lib/_shape_base_impl.py b/numpy/lib/_shape_base_impl.py
@@ -50,12 +50,12 @@ def _make_along_axis_idx(arr_shape, indices, axis):
     return tuple(fancy_index)
 
 
-def _take_along_axis_dispatcher(arr, indices, axis):
+def _take_along_axis_dispatcher(arr, indices, axis=None):
     return (arr, indices)
 
 
 @array_function_dispatch(_take_along_axis_dispatcher)
-def take_along_axis(arr, indices, axis):
+def take_along_axis(arr, indices, axis=-1):
     """
     Take values from the input array by matching 1d index and data slices.
 
@@ -71,14 +71,17 @@ def take_along_axis(arr, indices, axis):
     arr : ndarray (Ni..., M, Nk...)
         Source array
     indices : ndarray (Ni..., J, Nk...)
-        Indices to take along each 1d slice of `arr`. This must match the
-        dimension of arr, but dimensions Ni and Nj only need to broadcast
-        against `arr`.
-    axis : int
+        Indices to take along each 1d slice of ``arr``. This must match the
+        dimension of ``arr``, but dimensions Ni and Nj only need to broadcast
+        against ``arr``.
+    axis : int or None, optional
         The axis to take 1d slices along. If axis is None, the input array is
         treated as if it had first been flattened to 1d, for consistency with
         `sort` and `argsort`.
 
+        .. versionchanged:: 2.3
+            The default value is now ``-1``.
+
     Returns
     -------
     out: ndarray (Ni..., J, Nk...)
diff --git a/numpy/lib/_shape_base_impl.pyi b/numpy/lib/_shape_base_impl.pyi
@@ -70,7 +70,7 @@ class _SupportsArrayWrap(Protocol):
 def take_along_axis(
     arr: _ScalarT | NDArray[_ScalarT],
     indices: NDArray[integer],
-    axis: int | None,
+    axis: int | None = ...,
 ) -> NDArray[_ScalarT]: ...
 
 def put_along_axis(
