diff options
author | czgdp1807 <gdp.1807@gmail.com> | 2021-09-03 12:17:26 +0530 |
---|---|---|
committer | czgdp1807 <gdp.1807@gmail.com> | 2021-09-03 12:17:26 +0530 |
commit | 781d0a7ac61ce007e65abcd4e30f2181e729ae61 (patch) | |
tree | f45f38a246bcefbca9ca8a08bd8ba55cbc6cdb15 /numpy/array_api/_manipulation_functions.py | |
parent | b341e4c3249817d2e14ddf71aa850a8a896b9303 (diff) | |
parent | 2ae1e068710174dc57b5ba5ad688517608efcf26 (diff) | |
download | numpy-781d0a7ac61ce007e65abcd4e30f2181e729ae61.tar.gz |
resolved conflicts
Diffstat (limited to 'numpy/array_api/_manipulation_functions.py')
-rw-r--r-- | numpy/array_api/_manipulation_functions.py | 86 |
1 files changed, 86 insertions, 0 deletions
diff --git a/numpy/array_api/_manipulation_functions.py b/numpy/array_api/_manipulation_functions.py new file mode 100644 index 000000000..c11866261 --- /dev/null +++ b/numpy/array_api/_manipulation_functions.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from ._array_object import Array +from ._data_type_functions import result_type + +from typing import List, Optional, Tuple, Union + +import numpy as np + +# Note: the function name is different here +def concat( + arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: Optional[int] = 0 +) -> Array: + """ + Array API compatible wrapper for :py:func:`np.concatenate <numpy.concatenate>`. + + See its docstring for more information. + """ + # Note: Casting rules here are different from the np.concatenate default + # (no for scalars with axis=None, no cross-kind casting) + dtype = result_type(*arrays) + arrays = tuple(a._array for a in arrays) + return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype)) + + +def expand_dims(x: Array, /, *, axis: int) -> Array: + """ + Array API compatible wrapper for :py:func:`np.expand_dims <numpy.expand_dims>`. + + See its docstring for more information. + """ + return Array._new(np.expand_dims(x._array, axis)) + + +def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array: + """ + Array API compatible wrapper for :py:func:`np.flip <numpy.flip>`. + + See its docstring for more information. + """ + return Array._new(np.flip(x._array, axis=axis)) + + +def reshape(x: Array, /, shape: Tuple[int, ...]) -> Array: + """ + Array API compatible wrapper for :py:func:`np.reshape <numpy.reshape>`. + + See its docstring for more information. + """ + return Array._new(np.reshape(x._array, shape)) + + +def roll( + x: Array, + /, + shift: Union[int, Tuple[int, ...]], + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, +) -> Array: + """ + Array API compatible wrapper for :py:func:`np.roll <numpy.roll>`. + + See its docstring for more information. + """ + return Array._new(np.roll(x._array, shift, axis=axis)) + + +def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array: + """ + Array API compatible wrapper for :py:func:`np.squeeze <numpy.squeeze>`. + + See its docstring for more information. + """ + return Array._new(np.squeeze(x._array, axis=axis)) + + +def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array: + """ + Array API compatible wrapper for :py:func:`np.stack <numpy.stack>`. + + See its docstring for more information. + """ + # Call result type here just to raise on disallowed type combinations + result_type(*arrays) + arrays = tuple(a._array for a in arrays) + return Array._new(np.stack(arrays, axis=axis)) |