Merge pull request #24242 from charris/backport-24187

charris · web-flow · commit 0409b768a243 · 2023-07-23T21:24:32.000-06:00
BUG: Fix the signature for np.array_api.take
diff --git a/numpy/array_api/_indexing_functions.py b/numpy/array_api/_indexing_functions.py
@@ -5,14 +5,16 @@
 
 import numpy as np
 
-def take(x: Array, indices: Array, /, *, axis: int) -> Array:
+def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.take <numpy.take>`.
 
     See its docstring for more information.
-    """    
+    """
+    if axis is None and x.ndim != 1:
+        raise ValueError("axis must be specified when ndim > 1")
     if indices.dtype not in _integer_dtypes:
         raise TypeError("Only integer dtypes are allowed in indexing")
-    if indices.ndim != 1: 
+    if indices.ndim != 1:
         raise ValueError("Only 1-dim indices array is supported")
     return Array._new(np.take(x._array, indices._array, axis=axis))
