Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit4a36606

Browse files
committed
API: Add device and to_device to ndarray
1 parent47b2a9b commit4a36606

File tree

13 files changed

+202
-22
lines changed

13 files changed

+202
-22
lines changed
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+
defto_device(self,device:L["cpu"],/,*,stream:None|int|Any= ...)->NDArray[Any]: ...
2465+
2466+
@property
2467+
defdevice(self)->L["cpu"]: ...
2468+
24642469
defbitwise_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
deflinspace(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+
ifdevicenotin ["cpu",None]:
128+
raiseValueError(
129+
f"Unsupported device:{device}. Only\"cpu\" is allowed."
130+
)
131+
122132
num=operator.index(num)
123133
ifnum<0:
124134
raiseValueError(

‎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
deflinspace(
@@ -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
deflinspace(
@@ -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
deflinspace(
@@ -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
deflinspace(
@@ -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
deflinspace(
@@ -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
deflinspace(
@@ -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
deflinspace(
@@ -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-
defempty_like(prototype,dtype=None,order=None,subok=None,shape=None):
84+
defempty_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
defempty_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
defempty_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
defempty_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
defempty_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-
defones(shape,dtype=None,order='C',*,like=None):
135+
defones(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+
ifdevicenotin ["cpu",None]:
189+
raiseValueError(
190+
f"Unsupported device:{device}. Only\"cpu\" is allowed."
191+
)
192+
185193
iflikeisnotNone:
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
returnres
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-
deffull(shape,fill_value,dtype=None,order='C',*,like=None):
283+
deffull(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+
ifdevicenotin ["cpu",None]:
334+
raiseValueError(
335+
f"Unsupported device:{device}. Only\"cpu\" is allowed."
336+
)
337+
319338
iflikeisnotNone:
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-
deffull_like(a,fill_value,dtype=None,order='K',subok=True,shape=None):
361+
deffull_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+
ifdevicenotin ["cpu",None]:
433+
raiseValueError(
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
returnres

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp