8000 API: Add device and to_device to ndarray · numpy/numpy@91b8d6f · GitHub
[go: up one dir, main page]
Skip to content

Commit 91b8d6f

Browse files
mtsokolmtsokol
committed
API: Add device and to_device to ndarray
1 parent 35c4319 commit 91b8d6f

22 files changed

+384
-51
lines changed

‎doc/release/upcoming_changes/25233.new_feature.rst

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
``ndarray.device`` and ``ndarray.to_device``
2+
--------------------------------------------
3+
4+
``ndarray.device`` attribute and ``ndarray.to_device`` function were
5+
added to `numpy.ndarray` class for Array API compatibility.
6+
7+
Additionally, ``device`` keyword-only arguments were added to:
8+
`numpy.asarray`, `numpy.arange`, `numpy.empty`, `numpy.empty_like`,
9+
`numpy.eye`, `numpy.full`, `numpy.full_like`, `numpy.linspace`,
10+
`numpy.ones`, `numpy.ones_like`, `numpy.zeros`, and `numpy.zeros_like`.

‎doc/source/reference/array_api.rst

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -517,12 +517,6 @@ Creation functions differences
517517
* - ``copy`` keyword argument to ``asarray``
518518
- **Compatible**
519519
-
520-
* - New ``device`` keyword argument to all array creation functions
521-
(``asarray``, ``arange``, ``empty``, ``empty_like``, ``eye``, ``full``,
522-
``full_like``, ``linspace``, ``ones``, ``ones_like``, ``zeros``, and
523-
``zeros_like``).
524-
- **Compatible**
525-
- ``device`` would effectively do nothing, since NumPy is CPU only.
526520

527521
Elementwise functions differences
528522
---------------------------------

‎numpy/__init__.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2470,6 +2470,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
24702470

24712471
def __array_namespace__(self, *, api_version: str = ...) -> Any: ...
24722472

2473+
def to_device(self, device: L["cpu"], /, *, stream: None | int | Any = ...) -> NDArray[Any]: ...
2474+
2475+
@property
2476+
def device(self) -> L["cpu"]: ...
2477+
24732478
def bitwise_count(
24742479
self,
24752480
out: None | NDArray[Any] = ...,

‎numpy/_core/_add_newdocs.py

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -912,7 +912,7 @@
912912

913913
add_newdoc('numpy._core.multiarray', 'asarray',
914914
"""
915-
asarray(a, dtype=None, order=None, *, like=None)
915+
asarray(a, dtype=None, order=None, *, device=None, like=None)
916916
917917
Convert the input to an array.
918918
@@ -931,6 +931,14 @@
931931
'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise
932932
'K' (keep) preserve input order
933933
Defaults to 'K'.
934+
device : str, optional
935+
The device on which to place the created array. Default: None.
936+
937+
.. note:
938+
939+
Only the ``"cpu"`` device is supported by NumPy.
940+
941+
.. versionadded:: 2.0.0
934942
${ARRAY_FUNCTION_LIKE}
935943
936944
.. versionadded:: 1.20.0
@@ -1184,7 +1192,7 @@
11841192

11851193
add_newdoc('numpy._core.multiarray', 'empty',
11861194
"""
1187-
empty(shape, dtype=float, order='C', *, like=None)
1195+
empty(shape, dtype=float, order='C', *, device=None, like=None)
11881196
11891197
Return a new array of given shape and type, without initializing entries.
11901198
@@ -1199,6 +1207,14 @@
11991207
Whether to store multi-dimensional data in row-major
12001208
(C-style) or column-major (Fortran-style) order in
12011209
memory.
1210+
device : str, optional
1211+
The device on which to place the created array. Default: None.
1212+
1213+
.. note:
1214+
1215+
Only the ``"cpu"`` device is supported by NumPy.
1216+
1217+
.. versionadded:: 2.0.0
12021218
${ARRAY_FUNCTION_LIKE}
12031219
12041220
.. versionadded:: 1.20.0
@@ -1676,7 +1692,7 @@
16761692

16771693
add_newdoc('numpy._core.multiarray', 'arange',
16781694
"""
1679-
arange([start,] stop[, step,], dtype=None, *, like=None)
1695+
arange([start,] stop[, step,], dtype=None, *, device=None, like=None)
16801696
16811697
Return evenly spaced values within a given interval.
16821698
@@ -1717,6 +1733,14 @@
17171733
dtype : dtype, optional
17181734
The type of the output array. If `dtype` is not given, infer the data
17191735
type from the other input arguments.
1736+
device : str, optional
1737+
The device on which to place the created array. Default: None.
1738+
1739+
.. note:
1740+
1741+
Only the ``"cpu"`` device is supported by NumPy.
1742+
1743+
.. versionadded:: 2.0.0
17201744
${ARRAY_FUNCTION_LIKE}
17211745
17221746
.. versionadded:: 1.20.0

‎numpy/_core/function_base.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,13 @@
1717

1818

1919
def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None,
20-
dtype=None, axis=None):
20+
dtype=None, axis=None, *, device=None):
2121
return (start, stop)
2222

2323

2424
@array_function_dispatch(_linspace_dispatcher)
2525
def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
26-
axis=0):
26+
axis=0, *, device=None):
2727
"""
2828
Return evenly spaced numbers over a specified interval.
2929
@@ -64,13 +64,20 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
6464
array of integers.
6565
6666
.. versionadded:: 1.9.0
67-
6867
axis : int, optional
6968
The axis in the result to store the samples. Relevant only if start
7069
or stop are array-like. By default (0), the samples will be along a
7170
new axis inserted at the beginning. Use -1 to get an axis at the end.
7271
7372
.. versionadded:: 1.16.0
73+
device : str, optional
74+
The device on which to place the created array. Default: None.
75+
76+
.. note:
77+
78+
Only the ``"cpu"`` device is supported by NumPy.
79+
80+
.. versionadded:: 2.0.0
7481
7582
Returns
7683
-------
@@ -119,6 +126,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
119126
>>> plt.show()
120127
121128
"""
129+
if device not in ["cpu", None]:
130+
raise ValueError(
131+
f"Unsupported device: {device}. Only \"cpu\" is allowed."
132+
)
133+
122134
num = operator.index(num)
123135
if num < 0:
124136
raise ValueError(

‎numpy/_core/function_base.pyi

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ def linspace(
2828
retstep: L[False] = ...,
2929
dtype: None = ...,
3030
axis: SupportsIndex = ...,
31+
*,
32+
device: None | L["cpu"] = ...,
3133
) -> NDArray[floating[Any]]: ...
3234
@overload
3335
def linspace(
@@ -38,6 +40,8 @@ def linspace(
3840
retstep: L[False] = ...,
3941
dtype: None = ...,
4042
axis: SupportsIndex = ...,
43+
*,
44+
device: None | L["cpu"] = ...,
4145
) -> NDArray[complexfloating[Any, Any]]: ...
4246
@overload
4347
def linspace(
@@ -48,6 +52,8 @@ def linspace(
4852
retstep: L[False] = ...,
4953
dtype: _DTypeLike[_SCT] = ...,
5054
axis: SupportsIndex = ...,
55+
*,
56+
device: None | L["cpu"] = ...,
5157
) -> NDArray[_SCT]: ...
5258
@overload
5359
def linspace(
@@ -58,6 +64,8 @@ def linspace(
5864
retstep: L[False] = ...,
5965
dtype: DTypeLike = ...,
6066
axis: SupportsIndex = ...,
67+
*,
68+
device: None | L["cpu"] = ...,
6169
) -> NDArray[Any]: ...
6270
@overload
6371
def linspace(
@@ -68,6 +76,8 @@ def linspace(
6876
retstep: L[True] = ...,
6977
dtype: None = ...,
7078
axis: SupportsIndex = ...,
79+
*,
80+
device: None | L["cpu"] = ...,
7181
) -> tuple[NDArray[floating[Any]], floating[Any]]: ...
7282
@overload
7383
def linspace(
@@ -78,6 +88,8 @@ def linspace(
7888
retstep: L[True] = ...,
7989
dtype: None = ...,
8090
axis: SupportsIndex = ...,
91+
*,
92+
device: None | L["cpu"] = ...,
8193
) -> tuple[NDArray[complexfloating[Any, Any]], complexfloating[Any, Any]]: ...
8294
@overload
8395
def linspace(
@@ -88,6 +100,8 @@ def linspace(
88100
retstep: L[True] = ...,
89101
dtype: _DTypeLike[_SCT] = ...,
90102
axis: SupportsIndex = ...,
103+
*,
104+
device: None | L["cpu"] = ...,
91105
) -> tuple[NDArray[_SCT], _SCT]: ...
92106
@overload
93107
def linspace(
@@ -98,6 +112,8 @@ def linspace(
98112
retstep: L[True] = ...,
99113
dtype: DTypeLike = ...,
100114
axis: SupportsIndex = ...,
115+
*,
116+
device: None | L["cpu"] = ...,
101117
) -> tuple[NDArray[Any], Any]: ...
102118

103119
@overload

‎numpy/_core/include/numpy/ndarraytypes.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,11 @@ typedef enum {
300300
NPY_BUSDAY_RAISE
301301
} NPY_BUSDAY_ROLL;
302302

303+
/* Device enum for Array API compatibility */
304+
typedef enum {
305+
NPY_DEVICE_CPU = 0,
306+
} NPY_DEVICE;
307+
303308
/************************************************************
304309
* NumPy Auxiliary Data for inner loops, sort functions, etc.
305310
************************************************************/

‎numpy/_core/multiarray.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,12 @@
8181

8282

8383
@array_function_from_c_func_and_dispatcher(_multiarray_umath.empty_like)
84-
def empty_like(prototype, dtype=None, order=None, subok=None, shape=None):
84+
def empty_like(
85+
prototype, dtype=None, order=None, subok=None, shape=None, *, device=None
86+
):
8587
"""
86-
empty_like(prototype, dtype=None, order='K', subok=True, shape=None)
88+
empty_like(prototype, dtype=None, order='K', subok=True, shape=None,
89+
shape=None, *, device=None)
8790
8891
Return a new array with the same shape and type as a given array.
8992
@@ -113,6 +116,14 @@ def empty_like(prototype, dtype=None, order=None, subok=None, shape=None):
113116
order='C' is implied.
114117
115118
.. versionadded:: 1.17.0
119+
device : str, optional
120+
The device on which to place the created array. Default: None.
121+
122+
.. note:
123+
124+
Only the ``"cpu"`` device is supported by NumPy.
125+
126+
.. versionadded:: 2.0.0
116127
117128
Returns
118129
-------

0 commit comments

Comments
 (0)
0