| from __future__ import annotations | |
| from ._array_object import Array | |
| from ._dtypes import _integer_dtypes | |
| import numpy as np | |
| 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: | |
| raise ValueError("Only 1-dim indices array is supported") | |
| return Array._new(np.take(x._array, indices._array, axis=axis)) | |