From 26cd48fd452ea4f393aaa07e1d94983cbdcc9856 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= Date: Thu, 23 Nov 2023 13:49:13 +0100 Subject: [PATCH 1/3] API: Add device and to_device to ndarray --- .../upcoming_changes/25233.new_feature.rst | 10 ++ doc/source/reference/array_api.rst | 6 -- numpy/__init__.pyi | 5 + numpy/_core/_add_newdocs.py | 30 +++++- numpy/_core/function_base.py | 18 +++- numpy/_core/function_base.pyi | 16 ++++ numpy/_core/include/numpy/ndarraytypes.h | 5 + numpy/_core/multiarray.py | 15 ++- numpy/_core/multiarray.pyi | 31 ++++++ numpy/_core/numeric.py | 96 ++++++++++++++++--- numpy/_core/numeric.pyi | 60 +++++++++--- numpy/_core/src/multiarray/conversion_utils.c | 22 +++++ numpy/_core/src/multiarray/conversion_utils.h | 6 ++ numpy/_core/src/multiarray/getset.c | 10 ++ numpy/_core/src/multiarray/methods.c | 33 +++++++ numpy/_core/src/multiarray/multiarraymodule.c | 13 +++ numpy/_core/src/multiarray/multiarraymodule.h | 1 + numpy/_core/tests/test_multiarray.py | 26 +++++ numpy/lib/_twodim_base_impl.py | 15 ++- numpy/lib/_twodim_base_impl.pyi | 4 + numpy/matrixlib/tests/test_defmatrix.py | 2 +- tools/ci/array-api-skips.txt | 13 --- 22 files changed, 384 insertions(+), 53 deletions(-) create mode 100644 doc/release/upcoming_changes/25233.new_feature.rst diff --git a/doc/release/upcoming_changes/25233.new_feature.rst b/doc/release/upcoming_changes/25233.new_feature.rst new file mode 100644 index 000000000000..5e0ed5898a0c --- /dev/null +++ b/doc/release/upcoming_changes/25233.new_feature.rst @@ -0,0 +1,10 @@ +``ndarray.device`` and ``ndarray.to_device`` +-------------------------------------------- + +``ndarray.device`` attribute and ``ndarray.to_device`` function were +added to `numpy.ndarray` class for Array API compatibility. + +Additionally, ``device`` keyword-only arguments were added to: +`numpy.asarray`, `numpy.arange`, `numpy.empty`, `numpy.empty_like`, +`numpy.eye`, `numpy.full`, `numpy.full_like`, `numpy.linspace`, +`numpy.ones`, `numpy.ones_like`, `numpy.zeros`, and `numpy.zeros_like`. diff --git a/doc/source/reference/array_api.rst b/doc/source/reference/array_api.rst index 33346b9de4c9..87430221bad0 100644 --- a/doc/source/reference/array_api.rst +++ b/doc/source/reference/array_api.rst @@ -501,12 +501,6 @@ Creation functions differences * - ``copy`` keyword argument to ``asarray`` - **Compatible** - - * - New ``device`` keyword argument to all array creation functions - (``asarray``, ``arange``, ``empty``, ``empty_like``, ``eye``, ``full``, - ``full_like``, ``linspace``, ``ones``, ``ones_like``, ``zeros``, and - ``zeros_like``). - - **Compatible** - - ``device`` would effectively do nothing, since NumPy is CPU only. Elementwise functions differences --------------------------------- diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 2e28f2a1c297..a0728322b3b7 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -2550,6 +2550,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): def __array_namespace__(self, *, api_version: str = ...) -> Any: ... + def to_device(self, device: L["cpu"], /, *, stream: None | int | Any = ...) -> NDArray[Any]: ... + + @property + def device(self) -> L["cpu"]: ... + def bitwise_count( self, out: None | NDArray[Any] = ..., diff --git a/numpy/_core/_add_newdocs.py b/numpy/_core/_add_newdocs.py index 045c29d0eefc..95d87326c05a 100644 --- a/numpy/_core/_add_newdocs.py +++ b/numpy/_core/_add_newdocs.py @@ -912,7 +912,7 @@ add_newdoc('numpy._core.multiarray', 'asarray', """ - asarray(a, dtype=None, order=None, *, like=None) + asarray(a, dtype=None, order=None, *, device=None, like=None) Convert the input to an array. @@ -931,6 +931,14 @@ 'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise 'K' (keep) preserve input order Defaults to 'K'. + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 @@ -1184,7 +1192,7 @@ add_newdoc('numpy._core.multiarray', 'empty', """ - empty(shape, dtype=float, order='C', *, like=None) + empty(shape, dtype=float, order='C', *, device=None, like=None) Return a new array of given shape and type, without initializing entries. @@ -1199,6 +1207,14 @@ Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 @@ -1676,7 +1692,7 @@ add_newdoc('numpy._core.multiarray', 'arange', """ - arange([start,] stop[, step,], dtype=None, *, like=None) + arange([start,] stop[, step,], dtype=None, *, device=None, like=None) Return evenly spaced values within a given interval. @@ -1717,6 +1733,14 @@ dtype : dtype, optional The type of the output array. If `dtype` is not given, infer the data type from the other input arguments. + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 diff --git a/numpy/_core/function_base.py b/numpy/_core/function_base.py index 27cc88729fe6..8547f310a93c 100644 --- a/numpy/_core/function_base.py +++ b/numpy/_core/function_base.py @@ -17,13 +17,13 @@ def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None, - dtype=None, axis=None): + dtype=None, axis=None, *, device=None): return (start, stop) @array_function_dispatch(_linspace_dispatcher) def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, - axis=0): + axis=0, *, device=None): """ Return evenly spaced numbers over a specified interval. @@ -64,13 +64,20 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, array of integers. .. versionadded:: 1.9.0 - axis : int, optional The axis in the result to store the samples. Relevant only if start or stop are array-like. By default (0), the samples will be along a new axis inserted at the beginning. Use -1 to get an axis at the end. .. versionadded:: 1.16.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -119,6 +126,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, >>> plt.show() """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) + num = operator.index(num) if num < 0: raise ValueError( diff --git a/numpy/_core/function_base.pyi b/numpy/_core/function_base.pyi index 73e9706a2b9d..59c3d6b4ea2c 100644 --- a/numpy/_core/function_base.pyi +++ b/numpy/_core/function_base.pyi @@ -28,6 +28,8 @@ def linspace( retstep: L[False] = ..., dtype: None = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[floating[Any]]: ... @overload def linspace( @@ -38,6 +40,8 @@ def linspace( retstep: L[False] = ..., dtype: None = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[complexfloating[Any, Any]]: ... @overload def linspace( @@ -48,6 +52,8 @@ def linspace( retstep: L[False] = ..., dtype: _DTypeLike[_SCT] = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def linspace( @@ -58,6 +64,8 @@ def linspace( retstep: L[False] = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload def linspace( @@ -68,6 +76,8 @@ def linspace( retstep: L[True] = ..., dtype: None = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> tuple[NDArray[floating[Any]], floating[Any]]: ... @overload def linspace( @@ -78,6 +88,8 @@ def linspace( retstep: L[True] = ..., dtype: None = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> tuple[NDArray[complexfloating[Any, Any]], complexfloating[Any, Any]]: ... @overload def linspace( @@ -88,6 +100,8 @@ def linspace( retstep: L[True] = ..., dtype: _DTypeLike[_SCT] = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> tuple[NDArray[_SCT], _SCT]: ... @overload def linspace( @@ -98,6 +112,8 @@ def linspace( retstep: L[True] = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., + *, + device: None | L["cpu"] = ..., ) -> tuple[NDArray[Any], Any]: ... @overload diff --git a/numpy/_core/include/numpy/ndarraytypes.h b/numpy/_core/include/numpy/ndarraytypes.h index b5169d5cf969..2ff203fd3ec8 100644 --- a/numpy/_core/include/numpy/ndarraytypes.h +++ b/numpy/_core/include/numpy/ndarraytypes.h @@ -300,6 +300,11 @@ typedef enum { NPY_BUSDAY_RAISE } NPY_BUSDAY_ROLL; +/* Device enum for Array API compatibility */ +typedef enum { + NPY_DEVICE_CPU = 0, +} NPY_DEVICE; + /************************************************************ * NumPy Auxiliary Data for inner loops, sort functions, etc. ************************************************************/ diff --git a/numpy/_core/multiarray.py b/numpy/_core/multiarray.py index 55161def4807..677818ff3ff5 100644 --- a/numpy/_core/multiarray.py +++ b/numpy/_core/multiarray.py @@ -81,9 +81,12 @@ @array_function_from_c_func_and_dispatcher(_multiarray_umath.empty_like) -def empty_like(prototype, dtype=None, order=None, subok=None, shape=None): +def empty_like( + prototype, dtype=None, order=None, subok=None, shape=None, *, device=None +): """ - empty_like(prototype, dtype=None, order='K', subok=True, shape=None) + empty_like(prototype, dtype=None, order='K', subok=True, shape=None, + shape=None, *, device=None) Return a new array with the same shape and type as a given array. @@ -113,6 +116,14 @@ def empty_like(prototype, dtype=None, order=None, subok=None, shape=None): order='C' is implied. .. versionadded:: 1.17.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- diff --git a/numpy/_core/multiarray.pyi b/numpy/_core/multiarray.pyi index 7df68adb40a4..2618cdd9f495 100644 --- a/numpy/_core/multiarray.pyi +++ b/numpy/_core/multiarray.pyi @@ -132,6 +132,8 @@ def empty_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> _ArrayType: ... @overload def empty_like( @@ -140,6 +142,8 @@ def empty_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def empty_like( @@ -148,6 +152,8 @@ def empty_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload def empty_like( @@ -156,6 +162,8 @@ def empty_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def empty_like( @@ -164,6 +172,8 @@ def empty_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload @@ -228,6 +238,7 @@ def zeros( dtype: None = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload @@ -236,6 +247,7 @@ def zeros( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -244,6 +256,7 @@ def zeros( dtype: DTypeLike, order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @@ -253,6 +266,7 @@ def empty( dtype: None = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload @@ -261,6 +275,7 @@ def empty( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -269,6 +284,7 @@ def empty( dtype: DTypeLike, order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @@ -468,6 +484,7 @@ def asarray( dtype: None = ..., order: _OrderKACF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -476,6 +493,7 @@ def asarray( dtype: None = ..., order: _OrderKACF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -484,6 +502,7 @@ def asarray( dtype: _DTypeLike[_SCT], order: _OrderKACF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -492,6 +511,7 @@ def asarray( dtype: DTypeLike, order: _OrderKACF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @@ -714,6 +734,7 @@ def arange( # type: ignore[misc] stop: _IntLike_co, /, *, dtype: None = ..., + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[signedinteger[Any]]: ... @overload @@ -723,6 +744,7 @@ def arange( # type: ignore[misc] step: _IntLike_co = ..., dtype: None = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[signedinteger[Any]]: ... @overload @@ -730,6 +752,7 @@ def arange( # type: ignore[misc] stop: _FloatLike_co, /, *, dtype: None = ..., + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[floating[Any]]: ... @overload @@ -739,6 +762,7 @@ def arange( # type: ignore[misc] step: _FloatLike_co = ..., dtype: None = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[floating[Any]]: ... @overload @@ -746,6 +770,7 @@ def arange( stop: _TD64Like_co, /, *, dtype: None = ..., + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[timedelta64]: ... @overload @@ -755,6 +780,7 @@ def arange( step: _TD64Like_co = ..., dtype: None = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[timedelta64]: ... @overload @@ -764,6 +790,7 @@ def arange( # both start and stop must always be specified for datetime64 step: datetime64 = ..., dtype: None = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[datetime64]: ... @overload @@ -771,6 +798,7 @@ def arange( stop: Any, /, *, dtype: _DTypeLike[_SCT], + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -780,6 +808,7 @@ def arange( step: Any = ..., dtype: _DTypeLike[_SCT] = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -787,6 +816,7 @@ def arange( stop: Any, /, *, dtype: DTypeLike, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -796,6 +826,7 @@ def arange( step: Any = ..., dtype: DTypeLike = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index 6df13cc7aaf4..4bdd790c8efc 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -55,12 +55,16 @@ '_get_promotion_state', '_set_promotion_state'] -def _zeros_like_dispatcher(a, dtype=None, order=None, subok=None, shape=None): +def _zeros_like_dispatcher( + a, dtype=None, order=None, subok=None, shape=None, *, device=None +): return (a,) @array_function_dispatch(_zeros_like_dispatcher) -def zeros_like(a, dtype=None, order='K', subok=True, shape=None): +def zeros_like( + a, dtype=None, order='K', subok=True, shape=None, *, device=None +): """ Return an array of zeros with the same shape and type as a given array. @@ -90,6 +94,14 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None): order='C' is implied. .. versionadded:: 1.17.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -121,7 +133,9 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None): array([0., 0., 0.]) """ - res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape) + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) # needed instead of a 0 to get same result as zeros for string dtypes z = zeros(1, dtype=res.dtype) multiarray.copyto(res, z, casting='unsafe') @@ -130,7 +144,7 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None): @set_array_function_like_doc @set_module('numpy') -def ones(shape, dtype=None, order='C', *, like=None): +def ones(shape, dtype=None, order='C', *, device=None, like=None): """ Return a new array of given shape and type, filled with ones. @@ -145,6 +159,14 @@ def ones(shape, dtype=None, order='C', *, like=None): Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 @@ -161,7 +183,6 @@ def ones(shape, dtype=None, order='C', *, like=None): zeros : Return a new array setting values to zero. full : Return a new array of given shape filled with value. - Examples -------- >>> np.ones(5) @@ -180,6 +201,11 @@ def ones(shape, dtype=None, order='C', *, like=None): [1., 1.]]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) + if like is not None: return _ones_with_like(like, shape, dtype=dtype, order=order) @@ -191,12 +217,16 @@ def ones(shape, dtype=None, order='C', *, like=None): _ones_with_like = array_function_dispatch()(ones) -def _ones_like_dispatcher(a, dtype=None, order=None, subok=None, shape=None): +def _ones_like_dispatcher( + a, dtype=None, order=None, subok=None, shape=None, *, device=None +): return (a,) @array_function_dispatch(_ones_like_dispatcher) -def ones_like(a, dtype=None, order='K', subok=True, shape=None): +def ones_like( + a, dtype=None, order='K', subok=True, shape=None, *, device=None +): """ Return an array of ones with the same shape and type as a given array. @@ -226,6 +256,14 @@ def ones_like(a, dtype=None, order='K', subok=True, shape=None): order='C' is implied. .. versionadded:: 1.17.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -257,18 +295,22 @@ def ones_like(a, dtype=None, order='K', subok=True, shape=None): array([1., 1., 1.]) """ - res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape) + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) multiarray.copyto(res, 1, casting='unsafe') return res -def _full_dispatcher(shape, fill_value, dtype=None, order=None, *, like=None): +def _full_dispatcher( + shape, fill_value, dtype=None, order=None, *, device=None, like=None +): return(like,) @set_array_function_like_doc @set_module('numpy') -def full(shape, fill_value, dtype=None, order='C', *, like=None): +def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None): """ Return a new array of given shape and type, filled with `fill_value`. @@ -284,6 +326,14 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None): order : {'C', 'F'}, optional Whether to store multidimensional data in C- or Fortran-contiguous (row- or column-wise) order in memory. + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 @@ -314,6 +364,11 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None): [1, 2]]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) + if like is not None: return _full_with_like( like, shape, fill_value, dtype=dtype, order=order) @@ -330,13 +385,17 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None): def _full_like_dispatcher( - a, fill_value, dtype=None, order=None, subok=None, shape=None + a, fill_value, dtype=None, order=None, subok=None, shape=None, + *, device=None ): return (a,) @array_function_dispatch(_full_like_dispatcher) -def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None): +def full_like( + a, fill_value, dtype=None, order='K', subok=True, shape=None, + *, device=None +): """ Return a full array with the same shape and type as a given array. @@ -364,6 +423,14 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None): order='C' is implied. .. versionadded:: 1.17.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -400,6 +467,11 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None): [[ 0, 0, 255], [ 0, 0, 255]]]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) + res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape) multiarray.copyto(res, fill_value, casting='unsafe') return res diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi index 3c672d720a73..a24c368cbd08 100644 --- a/numpy/_core/numeric.pyi +++ b/numpy/_core/numeric.pyi @@ -3,7 +3,7 @@ from typing import ( Any, overload, TypeVar, - Literal, + Literal as L, SupportsAbs, SupportsIndex, NoReturn, @@ -53,7 +53,7 @@ _T = TypeVar("_T") _SCT = TypeVar("_SCT", bound=generic) _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) -_CorrelateMode = Literal["valid", "same", "full"] +_CorrelateMode = L["valid", "same", "full"] __all__: list[str] @@ -62,8 +62,10 @@ def zeros_like( a: _ArrayType, dtype: None = ..., order: _OrderKACF = ..., - subok: Literal[True] = ..., + subok: L[True] = ..., shape: None = ..., + *, + device: None | L["cpu"] = ..., ) -> _ArrayType: ... @overload def zeros_like( @@ -72,6 +74,8 @@ def zeros_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def zeros_like( @@ -80,6 +84,8 @@ def zeros_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload def zeros_like( @@ -88,6 +94,8 @@ def zeros_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def zeros_like( @@ -96,6 +104,8 @@ def zeros_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload @@ -104,6 +114,7 @@ def ones( dtype: None = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload @@ -112,6 +123,7 @@ def ones( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -120,6 +132,7 @@ def ones( dtype: DTypeLike, order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @@ -128,8 +141,10 @@ def ones_like( a: _ArrayType, dtype: None = ..., order: _OrderKACF = ..., - subok: Literal[True] = ..., + subok: L[True] = ..., shape: None = ..., + *, + device: None | L["cpu"] = ..., ) -> _ArrayType: ... @overload def ones_like( @@ -138,6 +153,8 @@ def ones_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def ones_like( @@ -146,6 +163,8 @@ def ones_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload def ones_like( @@ -154,6 +173,8 @@ def ones_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def ones_like( @@ -162,6 +183,8 @@ def ones_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload @@ -171,6 +194,7 @@ def full( dtype: None = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -180,6 +204,7 @@ def full( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -189,6 +214,7 @@ def full( dtype: DTypeLike, order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @@ -198,8 +224,10 @@ def full_like( fill_value: Any, dtype: None = ..., order: _OrderKACF = ..., - subok: Literal[True] = ..., + subok: L[True] = ..., shape: None = ..., + *, + device: None | L["cpu"] = ..., ) -> _ArrayType: ... @overload def full_like( @@ -209,6 +237,8 @@ def full_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike = ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def full_like( @@ -218,6 +248,8 @@ def full_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload def full_like( @@ -227,6 +259,8 @@ def full_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[_SCT]: ... @overload def full_like( @@ -236,6 +270,8 @@ def full_like( order: _OrderKACF = ..., subok: bool = ..., shape: None | _ShapeLike= ..., + *, + device: None | L["cpu"] = ..., ) -> NDArray[Any]: ... @overload @@ -243,7 +279,7 @@ def count_nonzero( a: ArrayLike, axis: None = ..., *, - keepdims: Literal[False] = ..., + keepdims: L[False] = ..., ) -> int: ... @overload def count_nonzero( @@ -587,37 +623,37 @@ def cross( def indices( dimensions: Sequence[int], dtype: type[int] = ..., - sparse: Literal[False] = ..., + sparse: L[False] = ..., ) -> NDArray[int_]: ... @overload def indices( dimensions: Sequence[int], dtype: type[int] = ..., - sparse: Literal[True] = ..., + sparse: L[True] = ..., ) -> tuple[NDArray[int_], ...]: ... @overload def indices( dimensions: Sequence[int], dtype: _DTypeLike[_SCT], - sparse: Literal[False] = ..., + sparse: L[False] = ..., ) -> NDArray[_SCT]: ... @overload def indices( dimensions: Sequence[int], dtype: _DTypeLike[_SCT], - sparse: Literal[True], + sparse: L[True], ) -> tuple[NDArray[_SCT], ...]: ... @overload def indices( dimensions: Sequence[int], dtype: DTypeLike, - sparse: Literal[False] = ..., + sparse: L[False] = ..., ) -> NDArray[Any]: ... @overload def indices( dimensions: Sequence[int], dtype: DTypeLike, - sparse: Literal[True], + sparse: L[True], ) -> tuple[NDArray[Any], ...]: ... def fromfunction( diff --git a/numpy/_core/src/multiarray/conversion_utils.c b/numpy/_core/src/multiarray/conversion_utils.c index 12ed45853e64..ba1d12810f61 100644 --- a/numpy/_core/src/multiarray/conversion_utils.c +++ b/numpy/_core/src/multiarray/conversion_utils.c @@ -18,6 +18,7 @@ #include "conversion_utils.h" #include "alloc.h" #include "npy_buffer.h" +#include "multiarraymodule.h" static int PyArray_PyIntAsInt_ErrMsg(PyObject *o, const char * msg) NPY_GCC_NONNULL(2); @@ -1352,3 +1353,24 @@ PyArray_IntTupleFromIntp(int len, npy_intp const *vals) fail: return intTuple; } + +/* + * Device string converter. + */ +NPY_NO_EXPORT int +PyArray_DeviceConverterOptional(PyObject *object, NPY_DEVICE *device) +{ + if (object == Py_None) { + return NPY_SUCCEED; + } + + if (PyUnicode_Check(object) && + PyUnicode_Compare(object, npy_ma_str_cpu) == 0) { + *device = NPY_DEVICE_CPU; + return NPY_SUCCEED; + } + + PyErr_SetString(PyExc_ValueError, + "Device not understood. Only \"cpu\" is allowed."); + return NPY_FAIL; +} diff --git a/numpy/_core/src/multiarray/conversion_utils.h b/numpy/_core/src/multiarray/conversion_utils.h index 62f54749ef0d..6fca3af2253f 100644 --- a/numpy/_core/src/multiarray/conversion_utils.h +++ b/numpy/_core/src/multiarray/conversion_utils.h @@ -82,6 +82,12 @@ PyArray_SelectkindConverter(PyObject *obj, NPY_SELECTKIND *selectkind); NPY_NO_EXPORT int PyArray_ConvertMultiAxis(PyObject *axis_in, int ndim, npy_bool *out_axis_flags); +/* + * Device string converter. + */ +NPY_NO_EXPORT int +PyArray_DeviceConverterOptional(PyObject *object, NPY_DEVICE *device); + /** * WARNING: This flag is a bad idea, but was the only way to both * 1) Support unpickling legacy pickles with object types. diff --git a/numpy/_core/src/multiarray/getset.c b/numpy/_core/src/multiarray/getset.c index 7200072d054f..6ab977a8822e 100644 --- a/numpy/_core/src/multiarray/getset.c +++ b/numpy/_core/src/multiarray/getset.c @@ -967,6 +967,12 @@ array_itemset(PyArrayObject *self, PyObject *args) return NULL; } +static PyObject * +array_device(PyArrayObject *self, void *NPY_UNUSED(ignored)) +{ + return PyUnicode_FromString("cpu"); +} + NPY_NO_EXPORT PyGetSetDef array_getsetlist[] = { {"ndim", (getter)array_ndim_get, @@ -1044,6 +1050,10 @@ NPY_NO_EXPORT PyGetSetDef array_getsetlist[] = { (getter)array_itemset, NULL, NULL, NULL}, + {"device", + (getter)array_device, + NULL, + NULL, NULL}, {"__array_interface__", (getter)array_interface_get, NULL, diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c index 0efc9311812f..cd7a902037df 100644 --- a/numpy/_core/src/multiarray/methods.c +++ b/numpy/_core/src/multiarray/methods.c @@ -2783,6 +2783,36 @@ array_array_namespace(PyArrayObject *self, PyObject *args, PyObject *kwds) return numpy_module; } +static PyObject * +array_to_device(PyArrayObject *self, PyObject *args, PyObject *kwds) +{ + static char *kwlist[] = {"", "stream", NULL}; + char *device = ""; + PyObject *stream = Py_None; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|$O:to_device", kwlist, + &device, + &stream)) { + return NULL; + } + + if (stream != Py_None) { + PyErr_SetString(PyExc_ValueError, + "The stream argument in to_device() " + "is not supported"); + return NULL; + } + + if (strcmp(device, "cpu") != 0) { + PyErr_Format(PyExc_ValueError, + "Unsupported device: %s.", device); + return NULL; + } + + Py_INCREF(self); + return (PyObject *)self; +} + NPY_NO_EXPORT PyMethodDef array_methods[] = { /* for subtypes */ @@ -3012,6 +3042,9 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = { {"__array_namespace__", (PyCFunction)array_array_namespace, METH_VARARGS | METH_KEYWORDS, NULL}, + {"to_device", + (PyCFunction)array_to_device, + METH_VARARGS | METH_KEYWORDS, NULL}, {NULL, NULL, 0, NULL} /* sentinel */ }; diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index 026b19446c8d..9d1f94d4ac5f 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1751,6 +1751,7 @@ array_asarray(PyObject *NPY_UNUSED(ignored), PyObject *op; npy_dtype_info dt_info = {NULL, NULL}; NPY_ORDER order = NPY_KEEPORDER; + NPY_DEVICE device = NPY_DEVICE_CPU; PyObject *like = Py_None; NPY_PREPARE_ARGPARSER; @@ -1759,6 +1760,7 @@ array_asarray(PyObject *NPY_UNUSED(ignored), "a", NULL, &op, "|dtype", &PyArray_DTypeOrDescrConverterOptional, &dt_info, "|order", &PyArray_OrderConverter, &order, + "$device", &PyArray_DeviceConverterOptional, &device, "$like", NULL, &like, NULL, NULL, NULL) < 0) { Py_XDECREF(dt_info.descr); @@ -1966,6 +1968,7 @@ array_empty(PyObject *NPY_UNUSED(ignored), NPY_ORDER order = NPY_CORDER; npy_bool is_f_order; PyArrayObject *ret = NULL; + NPY_DEVICE device = NPY_DEVICE_CPU; PyObject *like = Py_None; NPY_PREPARE_ARGPARSER; @@ -1973,6 +1976,7 @@ array_empty(PyObject *NPY_UNUSED(ignored), "shape", &PyArray_IntpConverter, &shape, "|dtype", &PyArray_DTypeOrDescrConverterOptional, &dt_info, "|order", &PyArray_OrderConverter, &order, + "$device", &PyArray_DeviceConverterOptional, &device, "$like", NULL, &like, NULL, NULL, NULL) < 0) { goto fail; @@ -2026,6 +2030,7 @@ array_empty_like(PyObject *NPY_UNUSED(ignored), int subok = 1; /* -1 is a special value meaning "not specified" */ PyArray_Dims shape = {NULL, -1}; + NPY_DEVICE device = NPY_DEVICE_CPU; NPY_PREPARE_ARGPARSER; @@ -2035,6 +2040,7 @@ array_empty_like(PyObject *NPY_UNUSED(ignored), "|order", &PyArray_OrderConverter, &order, "|subok", &PyArray_PythonPyIntFromInt, &subok, "|shape", &PyArray_OptionalIntpConverter, &shape, + "$device", &PyArray_DeviceConverterOptional, &device, NULL, NULL, NULL) < 0) { goto fail; } @@ -3055,6 +3061,7 @@ array_arange(PyObject *NPY_UNUSED(ignored), { PyObject *o_start = NULL, *o_stop = NULL, *o_step = NULL, *range=NULL; PyArray_Descr *typecode = NULL; + NPY_DEVICE device = NPY_DEVICE_CPU; PyObject *like = Py_None; NPY_PREPARE_ARGPARSER; @@ -3063,6 +3070,7 @@ array_arange(PyObject *NPY_UNUSED(ignored), "|stop", NULL, &o_stop, "|step", NULL, &o_step, "|dtype", &PyArray_DescrConverter2, &typecode, + "$device", &PyArray_DeviceConverterOptional, &device, "$like", NULL, &like, NULL, NULL, NULL) < 0) { Py_XDECREF(typecode); @@ -4774,6 +4782,7 @@ NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_axis2 = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_like = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_numpy = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_where = NULL; +NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_cpu = NULL; static int intern_strings(void) @@ -4834,6 +4843,10 @@ intern_strings(void) if (npy_ma_str_where == NULL) { return -1; } + npy_ma_str_cpu = PyUnicode_InternFromString("cpu"); + if (npy_ma_str_cpu == NULL) { + return -1; + } return 0; } diff --git a/numpy/_core/src/multiarray/multiarraymodule.h b/numpy/_core/src/multiarray/multiarraymodule.h index df505733bc63..2facfe76889f 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.h +++ b/numpy/_core/src/multiarray/multiarraymodule.h @@ -15,5 +15,6 @@ NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_axis2; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_like; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_numpy; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_where; +NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_cpu; #endif /* NUMPY_CORE_SRC_MULTIARRAY_MULTIARRAYMODULE_H_ */ diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py index d62902468e1d..9b4383512d27 100644 --- a/numpy/_core/tests/test_multiarray.py +++ b/numpy/_core/tests/test_multiarray.py @@ -10079,3 +10079,29 @@ def test_partition_fp(N, dtype): np.partition(arr, k, kind='introselect')) assert_arr_partitioned(np.sort(arr)[k], k, arr[np.argpartition(arr, k, kind='introselect')]) + + +class TestDevice: + """ + Test arr.device attribute and arr.to_device() method. + """ + def test_device(self): + arr = np.arange(5) + + assert arr.device == "cpu" + with assert_raises_regex( + AttributeError, + r"attribute 'device' of '(numpy.|)ndarray' objects is " + r"not writable" + ): + arr.device = "other" + + def test_to_device(self): + arr = np.arange(5) + + assert arr.to_device("cpu") is arr + with assert_raises_regex( + ValueError, + r"The stream argument in to_device\(\) is not supported" + ): + arr.to_device("cpu", stream=1) diff --git a/numpy/lib/_twodim_base_impl.py b/numpy/lib/_twodim_base_impl.py index b572146a3334..3160f80963f7 100644 --- a/numpy/lib/_twodim_base_impl.py +++ b/numpy/lib/_twodim_base_impl.py @@ -157,7 +157,7 @@ def flipud(m): @set_array_function_like_doc @set_module('numpy') -def eye(N, M=None, k=0, dtype=float, order='C', *, like=None): +def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None): """ Return a 2-D array with ones on the diagonal and zeros elsewhere. @@ -178,6 +178,14 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, like=None): column-major (Fortran-style) order in memory. .. versionadded:: 1.14.0 + device : str, optional + The device on which to place the created array. Default: None. + + .. note:: + + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} .. versionadded:: 1.20.0 @@ -204,6 +212,11 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, like=None): [0., 0., 0.]]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) + if like is not None: return _eye_with_like(like, N, M=M, k=k, dtype=dtype, order=order) if M is None: diff --git a/numpy/lib/_twodim_base_impl.pyi b/numpy/lib/_twodim_base_impl.pyi index ecc77c86a89a..4096976871d7 100644 --- a/numpy/lib/_twodim_base_impl.pyi +++ b/numpy/lib/_twodim_base_impl.pyi @@ -4,6 +4,7 @@ from typing import ( Any, overload, TypeVar, + Literal as L, ) import numpy as np @@ -64,6 +65,7 @@ def eye( dtype: None = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload @@ -74,6 +76,7 @@ def eye( dtype: _DTypeLike[_SCT] = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload @@ -84,6 +87,7 @@ def eye( dtype: DTypeLike = ..., order: _OrderCF = ..., *, + device: None | L["cpu"] = ..., like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... diff --git a/numpy/matrixlib/tests/test_defmatrix.py b/numpy/matrixlib/tests/test_defmatrix.py index 1aa88fccb5f6..81d955e86fa8 100644 --- a/numpy/matrixlib/tests/test_defmatrix.py +++ b/numpy/matrixlib/tests/test_defmatrix.py @@ -283,7 +283,7 @@ def test_instance_methods(self): 'argmin', 'choose', 'dump', 'dumps', 'fill', 'getfield', 'getA', 'getA1', 'item', 'nonzero', 'put', 'putmask', 'resize', 'searchsorted', 'setflags', 'setfield', 'sort', - 'partition', 'argpartition', 'newbyteorder', + 'partition', 'argpartition', 'newbyteorder', 'to_device', 'take', 'tofile', 'tolist', 'tostring', 'tobytes', 'all', 'any', 'sum', 'argmax', 'argmin', 'min', 'max', 'mean', 'var', 'ptp', 'prod', 'std', 'ctypes', 'itemset', 'bitwise_count', diff --git a/tools/ci/array-api-skips.txt b/tools/ci/array-api-skips.txt index a2bf766bb316..9364135bc8b0 100644 --- a/tools/ci/array-api-skips.txt +++ b/tools/ci/array-api-skips.txt @@ -39,29 +39,16 @@ array_api_tests/test_fft.py # finfo return type misalignment array_api_tests/test_data_type_functions.py::test_finfo[float32] -# missing names -array_api_tests/test_has_names.py::test_has_names[array_method-to_device] -array_api_tests/test_has_names.py::test_has_names[array_attribute-device] - # a few misalignments array_api_tests/test_operators_and_elementwise_functions.py array_api_tests/test_signatures.py::test_func_signature[std] array_api_tests/test_signatures.py::test_func_signature[var] array_api_tests/test_signatures.py::test_func_signature[asarray] -array_api_tests/test_signatures.py::test_func_signature[empty_like] -array_api_tests/test_signatures.py::test_func_signature[eye] -array_api_tests/test_signatures.py::test_func_signature[full] -array_api_tests/test_signatures.py::test_func_signature[full_like] -array_api_tests/test_signatures.py::test_func_signature[linspace] -array_api_tests/test_signatures.py::test_func_signature[ones] -array_api_tests/test_signatures.py::test_func_signature[ones_like] -array_api_tests/test_signatures.py::test_func_signature[zeros_like] array_api_tests/test_signatures.py::test_func_signature[reshape] array_api_tests/test_signatures.py::test_func_signature[argsort] array_api_tests/test_signatures.py::test_func_signature[sort] array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_rank] array_api_tests/test_signatures.py::test_array_method_signature[__array_namespace__] -array_api_tests/test_signatures.py::test_array_method_signature[to_device] # missing 'copy' keyword argument, 'newshape' should be named 'shape' array_api_tests/test_signatures.py::test_func_signature[reshape] From 6c510008a2515585d818e260c95343861c7cf761 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= Date: Wed, 10 Jan 2024 12:45:49 +0100 Subject: [PATCH 2/3] API: Add device keyword to fftfreq and rfftfreq Co-authored-by: Aaron Meurer --- numpy/fft/_helper.py | 26 ++++++++++++++++++++++++-- numpy/fft/_helper.pyi | 6 +++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/numpy/fft/_helper.py b/numpy/fft/_helper.py index b4c1dcb59b29..5739a3e1a164 100644 --- a/numpy/fft/_helper.py +++ b/numpy/fft/_helper.py @@ -121,7 +121,7 @@ def ifftshift(x, axes=None): @set_module('numpy.fft') -def fftfreq(n, d=1.0): +def fftfreq(n, d=1.0, device=None): """ Return the Discrete Fourier Transform sample frequencies. @@ -140,6 +140,13 @@ def fftfreq(n, d=1.0): Window length. d : scalar, optional Sample spacing (inverse of the sampling rate). Defaults to 1. + device : str, optional + The device on which to place the created array. Default: ``None``. + + .. note:: + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -157,6 +164,10 @@ def fftfreq(n, d=1.0): array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) if not isinstance(n, integer_types): raise ValueError("n should be an integer") val = 1.0 / (n * d) @@ -170,7 +181,7 @@ def fftfreq(n, d=1.0): @set_module('numpy.fft') -def rfftfreq(n, d=1.0): +def rfftfreq(n, d=1.0, device=None): """ Return the Discrete Fourier Transform sample frequencies (for usage with rfft, irfft). @@ -193,6 +204,13 @@ def rfftfreq(n, d=1.0): Window length. d : scalar, optional Sample spacing (inverse of the sampling rate). Defaults to 1. + device : str, optional + The device on which to place the created array. Default: ``None``. + + .. note:: + Only the ``"cpu"`` device is supported by NumPy. + + .. versionadded:: 2.0.0 Returns ------- @@ -213,6 +231,10 @@ def rfftfreq(n, d=1.0): array([ 0., 10., 20., 30., 40., 50.]) """ + if device not in ["cpu", None]: + raise ValueError( + f"Unsupported device: {device}. Only \"cpu\" is allowed." + ) if not isinstance(n, integer_types): raise ValueError("n should be an integer") val = 1.0/(n*d) diff --git a/numpy/fft/_helper.pyi b/numpy/fft/_helper.pyi index 9b65251900a3..a3c17fc675e7 100644 --- a/numpy/fft/_helper.pyi +++ b/numpy/fft/_helper.pyi @@ -1,4 +1,4 @@ -from typing import Any, TypeVar, overload +from typing import Any, TypeVar, overload, Literal as L from numpy import generic, integer, floating, complexfloating from numpy._typing import ( @@ -28,20 +28,24 @@ def ifftshift(x: ArrayLike, axes: None | _ShapeLike = ...) -> NDArray[Any]: ... def fftfreq( n: int | integer[Any], d: _ArrayLikeFloat_co = ..., + device: None | L["cpu"] = ..., ) -> NDArray[floating[Any]]: ... @overload def fftfreq( n: int | integer[Any], d: _ArrayLikeComplex_co = ..., + device: None | L["cpu"] = ..., ) -> NDArray[complexfloating[Any, Any]]: ... @overload def rfftfreq( n: int | integer[Any], d: _ArrayLikeFloat_co = ..., + device: None | L["cpu"] = ..., ) -> NDArray[floating[Any]]: ... @overload def rfftfreq( n: int | integer[Any], d: _ArrayLikeComplex_co = ..., + device: None | L["cpu"] = ..., ) -> NDArray[complexfloating[Any, Any]]: ... From 897f9d096328142a2f9d8bf2954c19d35ee64d46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= Date: Wed, 17 Jan 2024 13:25:43 +0100 Subject: [PATCH 3/3] Apply review comments --- .../upcoming_changes/25233.new_feature.rst | 4 +- numpy/_core/_add_newdocs.py | 19 ++----- numpy/_core/function_base.py | 14 ++--- numpy/_core/multiarray.py | 5 +- numpy/_core/numeric.py | 55 +++++-------------- numpy/_core/src/multiarray/conversion_utils.c | 5 +- numpy/_core/src/multiarray/multiarraymodule.c | 2 + numpy/fft/_helper.py | 24 ++------ numpy/lib/_twodim_base_impl.py | 16 ++---- 9 files changed, 44 insertions(+), 100 deletions(-) diff --git a/doc/release/upcoming_changes/25233.new_feature.rst b/doc/release/upcoming_changes/25233.new_feature.rst index 5e0ed5898a0c..712d93e3f17e 100644 --- a/doc/release/upcoming_changes/25233.new_feature.rst +++ b/doc/release/upcoming_changes/25233.new_feature.rst @@ -1,10 +1,12 @@ ``ndarray.device`` and ``ndarray.to_device`` -------------------------------------------- -``ndarray.device`` attribute and ``ndarray.to_device`` function were +``ndarray.device`` attribute and ``ndarray.to_device`` method were added to `numpy.ndarray` class for Array API compatibility. Additionally, ``device`` keyword-only arguments were added to: `numpy.asarray`, `numpy.arange`, `numpy.empty`, `numpy.empty_like`, `numpy.eye`, `numpy.full`, `numpy.full_like`, `numpy.linspace`, `numpy.ones`, `numpy.ones_like`, `numpy.zeros`, and `numpy.zeros_like`. + +For all these new arguments, only ``device="cpu"`` is supported. diff --git a/numpy/_core/_add_newdocs.py b/numpy/_core/_add_newdocs.py index 95d87326c05a..0a88da77f666 100644 --- a/numpy/_core/_add_newdocs.py +++ b/numpy/_core/_add_newdocs.py @@ -933,10 +933,7 @@ Defaults to 'K'. device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} @@ -946,8 +943,8 @@ Returns ------- out : ndarray - Array interpretation of `a`. No copy is performed if the input - is already an ndarray with matching dtype and order. If `a` is a + Array interpretation of ``a``. No copy is performed if the input + is already an ndarray with matching dtype and order. If ``a`` is a subclass of ndarray, a base class ndarray is returned. See Also @@ -1209,10 +1206,7 @@ memory. device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} @@ -1735,10 +1729,7 @@ type from the other input arguments. device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} diff --git a/numpy/_core/function_base.py b/numpy/_core/function_base.py index 8547f310a93c..4c2d0d6dc4fa 100644 --- a/numpy/_core/function_base.py +++ b/numpy/_core/function_base.py @@ -72,10 +72,7 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, .. versionadded:: 1.16.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -126,11 +123,6 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, >>> plt.show() """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) - num = operator.index(num) if num < 0: raise ValueError( @@ -152,7 +144,9 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, integer_dtype = _nx.issubdtype(dtype, _nx.integer) delta = stop - start - y = _nx.arange(0, num, dtype=dt).reshape((-1,) + (1,) * ndim(delta)) + y = _nx.arange( + 0, num, dtype=dt, device=device + ).reshape((-1,) + (1,) * ndim(delta)) # In-place multiplication y *= delta/div is faster, but prevents # the multiplicant from overriding what class is produced, and thus # prevents, e.g. use of Quantities, see gh-7142. Hence, we multiply diff --git a/numpy/_core/multiarray.py b/numpy/_core/multiarray.py index 677818ff3ff5..58e2dc114801 100644 --- a/numpy/_core/multiarray.py +++ b/numpy/_core/multiarray.py @@ -118,10 +118,7 @@ def empty_like( .. versionadded:: 1.17.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index 4bdd790c8efc..4e16f7e4688f 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -96,10 +96,7 @@ def zeros_like( .. versionadded:: 1.17.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -161,10 +158,7 @@ def ones(shape, dtype=None, order='C', *, device=None, like=None): memory. device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} @@ -201,15 +195,12 @@ def ones(shape, dtype=None, order='C', *, device=None, like=None): [1., 1.]]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) - if like is not None: - return _ones_with_like(like, shape, dtype=dtype, order=order) + return _ones_with_like( + like, shape, dtype=dtype, order=order, device=device + ) - a = empty(shape, dtype, order) + a = empty(shape, dtype, order, device=device) multiarray.copyto(a, 1, casting='unsafe') return a @@ -258,10 +249,7 @@ def ones_like( .. versionadded:: 1.17.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -328,10 +316,7 @@ def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None): (row- or column-wise) order in memory. device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} @@ -364,19 +349,15 @@ def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None): [1, 2]]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) - if like is not None: return _full_with_like( - like, shape, fill_value, dtype=dtype, order=order) + like, shape, fill_value, dtype=dtype, order=order, device=device + ) if dtype is None: fill_value = asarray(fill_value) dtype = fill_value.dtype - a = empty(shape, dtype, order) + a = empty(shape, dtype, order, device=device) multiarray.copyto(a, fill_value, casting='unsafe') return a @@ -425,10 +406,7 @@ def full_like( .. versionadded:: 1.17.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -467,12 +445,9 @@ def full_like( [[ 0, 0, 255], [ 0, 0, 255]]]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) - - res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape) + res = empty_like( + a, dtype=dtype, order=order, subok=subok, shape=shape, device=device + ) multiarray.copyto(res, fill_value, casting='unsafe') return res diff --git a/numpy/_core/src/multiarray/conversion_utils.c b/numpy/_core/src/multiarray/conversion_utils.c index ba1d12810f61..c72599ede7dc 100644 --- a/numpy/_core/src/multiarray/conversion_utils.c +++ b/numpy/_core/src/multiarray/conversion_utils.c @@ -1370,7 +1370,8 @@ PyArray_DeviceConverterOptional(PyObject *object, NPY_DEVICE *device) return NPY_SUCCEED; } - PyErr_SetString(PyExc_ValueError, - "Device not understood. Only \"cpu\" is allowed."); + PyErr_Format(PyExc_ValueError, + "Device not understood. Only \"cpu\" is allowed, " + "but received: %S", object); return NPY_FAIL; } diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index 9d1f94d4ac5f..11599070af2f 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -2181,6 +2181,7 @@ array_zeros(PyObject *NPY_UNUSED(ignored), NPY_ORDER order = NPY_CORDER; npy_bool is_f_order = NPY_FALSE; PyArrayObject *ret = NULL; + NPY_DEVICE device = NPY_DEVICE_CPU; PyObject *like = Py_None; NPY_PREPARE_ARGPARSER; @@ -2188,6 +2189,7 @@ array_zeros(PyObject *NPY_UNUSED(ignored), "shape", &PyArray_IntpConverter, &shape, "|dtype", &PyArray_DTypeOrDescrConverterOptional, &dt_info, "|order", &PyArray_OrderConverter, &order, + "$device", &PyArray_DeviceConverterOptional, &device, "$like", NULL, &like, NULL, NULL, NULL) < 0) { goto finish; diff --git a/numpy/fft/_helper.py b/numpy/fft/_helper.py index 5739a3e1a164..9f4512f90715 100644 --- a/numpy/fft/_helper.py +++ b/numpy/fft/_helper.py @@ -142,9 +142,7 @@ def fftfreq(n, d=1.0, device=None): Sample spacing (inverse of the sampling rate). Defaults to 1. device : str, optional The device on which to place the created array. Default: ``None``. - - .. note:: - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -164,18 +162,14 @@ def fftfreq(n, d=1.0, device=None): array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) if not isinstance(n, integer_types): raise ValueError("n should be an integer") val = 1.0 / (n * d) - results = empty(n, int) + results = empty(n, int, device=device) N = (n-1)//2 + 1 - p1 = arange(0, N, dtype=int) + p1 = arange(0, N, dtype=int, device=device) results[:N] = p1 - p2 = arange(-(n//2), 0, dtype=int) + p2 = arange(-(n//2), 0, dtype=int, device=device) results[N:] = p2 return results * val @@ -206,9 +200,7 @@ def rfftfreq(n, d=1.0, device=None): Sample spacing (inverse of the sampling rate). Defaults to 1. device : str, optional The device on which to place the created array. Default: ``None``. - - .. note:: - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 @@ -231,13 +223,9 @@ def rfftfreq(n, d=1.0, device=None): array([ 0., 10., 20., 30., 40., 50.]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) if not isinstance(n, integer_types): raise ValueError("n should be an integer") val = 1.0/(n*d) N = n//2 + 1 - results = arange(0, N, dtype=int) + results = arange(0, N, dtype=int, device=device) return results * val diff --git a/numpy/lib/_twodim_base_impl.py b/numpy/lib/_twodim_base_impl.py index 3160f80963f7..2a5a0b3917ec 100644 --- a/numpy/lib/_twodim_base_impl.py +++ b/numpy/lib/_twodim_base_impl.py @@ -180,10 +180,7 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None): .. versionadded:: 1.14.0 device : str, optional The device on which to place the created array. Default: None. - - .. note:: - - Only the ``"cpu"`` device is supported by NumPy. + For Array-API interoperability only, so must be ``"cpu"`` if passed. .. versionadded:: 2.0.0 ${ARRAY_FUNCTION_LIKE} @@ -212,16 +209,13 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None): [0., 0., 0.]]) """ - if device not in ["cpu", None]: - raise ValueError( - f"Unsupported device: {device}. Only \"cpu\" is allowed." - ) - if like is not None: - return _eye_with_like(like, N, M=M, k=k, dtype=dtype, order=order) + return _eye_with_like( + like, N, M=M, k=k, dtype=dtype, order=order, device=device + ) if M is None: M = N - m = zeros((N, M), dtype=dtype, order=order) + m = zeros((N, M), dtype=dtype, order=order, device=device) if k >= M: return m # Ensure M and k are integers, so we don't get any surprise casting