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

Commit 5037b48

Browse files
mtsokolmtsokol
committed
API: Add device and to_device to ndarray
1 parent 47b2a9b commit 5037b48

14 files changed

+203
-23
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
``ndarray.device`` and ``ndarray.to_device``
2+
--------------------------------------------
3+
4+
TODO

‎numpy/__init__.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2461,6 +2461,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
24612461
def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ...
24622462
def __dlpack_device__(self) -> tuple[int, L[0]]: ...
24632463

2464+
def to_device(self, device: L["cpu"], /, *, stream: None | int | Any = ...) -> NDArray[Any]: ...
2465+
2466+
@property
2467+
def device(self) -> L["cpu"]: ...
2468+
24642469
def bitwise_count(
24652470
self,
24662471
out: None | NDArray[Any] = ...,

‎numpy/_core/function_base.py

Lines changed: 12 additions & 2 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
@@ -72,6 +72,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
7272
7373
.. versionadded:: 1.16.0
7474
75+
device : str, optional
76+
The device on which to place the created array. Default: None.
77+
78+
.. versionadded:: 2.0.0
79+
7580
Returns
7681
-------
7782
samples : ndarray
@@ -119,6 +124,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
119124
>>> plt.show()
120125
121126
"""
127+
if device not in ["cpu", None]:
128+
raise ValueError(
129+
f"Unsupported device: {device}. Only \"cpu\" is allowed."
130+
)
131+
122132
num = operator.index(num)
123133
if num < 0:
124134
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/multiarray.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@
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
"""
8688
empty_like(prototype, dtype=None, order='K', subok=True, shape=None)
8789
@@ -113,6 +115,10 @@ def empty_like(prototype, dtype=None, order=None, subok=None, shape=None):
113115
order='C' is implied.
114116
115117
.. versionadded:: 1.17.0
118+
device : str, optional
119+
The device on which to place the created array. Default: None.
120+
121+
.. versionadded:: 2.0.0
116122
117123
Returns
118124
-------

‎numpy/_core/multiarray.pyi

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,8 @@ def empty_like(
132132
order: _OrderKACF = ...,
133133
subok: bool = ...,
134134
shape: None | _ShapeLike = ...,
135+
*,
136+
device: None | L["cpu"] = ...,
135137
) -> _ArrayType: ...
136138
@overload
137139
def empty_like(
@@ -140,6 +142,8 @@ def empty_like(
140142
order: _OrderKACF = ...,
141143
subok: bool = ...,
142144
shape: None | _ShapeLike = ...,
145+
*,
146+
device: None | L["cpu"] = ...,
143147
) -> NDArray[_SCT]: ...
144148
@overload
145149
def empty_like(
@@ -148,6 +152,8 @@ def empty_like(
148152
order: _OrderKACF = ...,
149153
subok: bool = ...,
150154
shape: None | _ShapeLike = ...,
155+
*,
156+
device: None | L["cpu"] = ...,
151157
) -> NDArray[Any]: ...
152158
@overload
153159
def empty_like(
@@ -156,6 +162,8 @@ def empty_like(
156162
order: _OrderKACF = ...,
157163
subok: bool = ...,
158164
shape: None | _ShapeLike = ...,
165+
*,
166+
device: None | L["cpu"] = ...,
159167
) -> NDArray[_SCT]: ...
160168
@overload
161169
def empty_like(
@@ -164,6 +172,8 @@ def empty_like(
164172
order: _OrderKACF = ...,
165173
subok: bool = ...,
166174
shape: None | _ShapeLike = ...,
175+
*,
176+
device: None | L["cpu"] = ...,
167177
) -> NDArray[Any]: ...
168178

169179
@overload

‎numpy/_core/numeric.py

Lines changed: 38 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None):
132132

133133
@set_array_function_like_doc
134134
@set_module('numpy')
135-
def ones(shape, dtype=None, order='C', *, like=None):
135+
def ones(shape, dtype=None, order='C', *, device=None, like=None):
136136
"""
137137
Return a new array of given shape and type, filled with ones.
138138
@@ -147,6 +147,10 @@ def ones(shape, dtype=None, order='C', *, like=None):
147147
Whether to store multi-dimensional data in row-major
148148
(C-style) or column-major (Fortran-style) order in
149149
memory.
150+
device : str, optional
151+
The device on which to place the created array. Default: None.
152+
153+
.. versionadded:: 2.0.0
150154
${ARRAY_FUNCTION_LIKE}
151155
152156
.. versionadded:: 1.20.0
@@ -163,7 +167,6 @@ def ones(shape, dtype=None, order='C', *, like=None):
163167
zeros : Return a new array setting values to zero.
164168
full : Return a new array of given shape filled with value.
165169
166-
167170
Examples
168171
--------
169172
>>> np.ones(5)
@@ -182,6 +185,11 @@ def ones(shape, dtype=None, order='C', *, like=None):
182185
[1., 1.]])
183186
184187
"""
188+
if device not in ["cpu", None]:
189+
raise ValueError(
190+
f"Unsupported device: {device}. Only \"cpu\" is allowed."
191+
)
192+
185193
if like is not None:
186194
return _ones_with_like(like, shape, dtype=dtype, order=order)
187195

@@ -264,13 +272,15 @@ def ones_like(a, dtype=None, order='K', subok=True, shape=None):
264272
return res
265273

266274

267-
def _full_dispatcher(shape, fill_value, dtype=None, order=None, *, like=None):
275+
def _full_dispatcher(
276+
shape, fill_value, dtype=None, order=None, *, device=None, like=None
277+
):
268278
return(like,)
269279

270280

271281
@set_array_function_like_doc
272282
@set_module('numpy')
273-
def full(shape, fill_value, dtype=None, order='C', *, like=None):
283+
def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None):
274284
"""
275285
Return a new array of given shape and type, filled with `fill_value`.
276286
@@ -286,6 +296,10 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None):
286296
order : {'C', 'F'}, optional
287297
Whether to store multidimensional data in C- or Fortran-contiguous
288298
(row- or column-wise) order in memory.
299+
device : str, optional
300+
The device on which to place the created array. Default: None.
301+
302+
.. versionadded:: 2.0.0
289303
${ARRAY_FUNCTION_LIKE}
290304
291305
.. versionadded:: 1.20.0
@@ -316,6 +330,11 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None):
316330
[1, 2]])
317331
318332
"""
333+
if device not in ["cpu", None]:
334+
raise ValueError(
335+
f"Unsupported device: {device}. Only \"cpu\" is allowed."
336+
)
337+
319338
if like is not None:
320339
return _full_with_like(
321340
like, shape, fill_value, dtype=dtype, order=order)
@@ -332,13 +351,17 @@ def full(shape, fill_value, dtype=None, order='C', *, like=None):
332351

333352

334353
def _full_like_dispatcher(
335-
a, fill_value, dtype=None, order=None, subok=None, shape=None
354+
a, fill_value, dtype=None, order=None, subok=None, shape=None,
355+
*, device=None
336356
):
337357
return (a,)
338358

339359

340360
@array_function_dispatch(_full_like_dispatcher)
341-
def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
361+
def full_like(
362+
a, fill_value, dtype=None, order='K', subok=True, shape=None,
363+
*, device=None
364+
):
342365
"""
343366
Return a full array with the same shape and type as a given array.
344367
@@ -366,6 +389,10 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
366389
order='C' is implied.
367390
368391
.. versionadded:: 1.17.0
392+
device : str, optional
393+
The device on which to place the created array. Default: None.
394+
395+
.. versionadded:: 2.0.0
369396
370397
Returns
371398
-------
@@ -402,6 +429,11 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
402429
[[ 0, 0, 255],
403430
[ 0, 0, 255]]])
404431
"""
432+
if device not in ["cpu", None]:
433+
raise ValueError(
434+
f"Unsupported device: {device}. Only \"cpu\" is allowed."
435+
)
436+
405437
res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
406438
multiarray.copyto(res, fill_value, casting='unsafe')
407439
return res

0 commit comments

Comments
 (0)
0