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

Commitebe3503

Browse files
committed
API: Add device and to_device to ndarray
1 parent06d7bdf commitebe3503

File tree

22 files changed

+339
-43
lines changed

22 files changed

+339
-43
lines changed
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
@@ -588,12 +588,6 @@ Creation functions differences
588588
* - ``copy`` keyword argument to ``asarray``
589589
- **Compatible**
590590
-
591-
* - New ``device`` keyword argument to all array creation functions
592-
(``asarray``, ``arange``, ``empty``, ``empty_like``, ``eye``, ``full``,
593-
``full_like``, ``linspace``, ``ones``, ``ones_like``, ``zeros``, and
594-
``zeros_like``).
595-
- **Compatible**
596-
- ``device`` would effectively do nothing, since NumPy is CPU only.
597591

598592
Elementwise functions differences
599593
---------------------------------

‎numpy/__init__.pyi‎

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

24662466
def__array_namespace__(self,*,api_version:str= ...)->Any: ...
24672467

2468+
defto_device(self,device:L["cpu"],/,*,stream:None|int|Any= ...)->NDArray[Any]: ...
2469+
2470+
@property
2471+
defdevice(self)->L["cpu"]: ...
2472+
24682473
defbitwise_count(
24692474
self,
24702475
out:None|NDArray[Any]= ...,

‎numpy/_core/_add_newdocs.py‎

Lines changed: 15 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,10 @@
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+
.. versionadded:: 2.0.0
934938
${ARRAY_FUNCTION_LIKE}
935939
936940
.. versionadded:: 1.20.0
@@ -1184,7 +1188,7 @@
11841188

11851189
add_newdoc('numpy._core.multiarray','empty',
11861190
"""
1187-
empty(shape, dtype=float, order='C', *, like=None)
1191+
empty(shape, dtype=float, order='C', *,device=None,like=None)
11881192
11891193
Return a new array of given shape and type, without initializing entries.
11901194
@@ -1199,6 +1203,10 @@
11991203
Whether to store multi-dimensional data in row-major
12001204
(C-style) or column-major (Fortran-style) order in
12011205
memory.
1206+
device : str, optional
1207+
The device on which to place the created array. Default: None.
1208+
1209+
.. versionadded:: 2.0.0
12021210
${ARRAY_FUNCTION_LIKE}
12031211
12041212
.. versionadded:: 1.20.0
@@ -1676,7 +1684,7 @@
16761684

16771685
add_newdoc('numpy._core.multiarray','arange',
16781686
"""
1679-
arange([start,] stop[, step,], dtype=None, *, like=None)
1687+
arange([start,] stop[, step,], dtype=None, *,device=None,like=None)
16801688
16811689
Return evenly spaced values within a given interval.
16821690
@@ -1717,6 +1725,10 @@
17171725
dtype : dtype, optional
17181726
The type of the output array. If `dtype` is not given, infer the data
17191727
type from the other input arguments.
1728+
device : str, optional
1729+
The device on which to place the created array. Default: None.
1730+
1731+
.. versionadded:: 2.0.0
17201732
${ARRAY_FUNCTION_LIKE}
17211733
17221734
.. versionadded:: 1.20.0

‎numpy/_core/function_base.py‎

Lines changed: 11 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
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
@@ -64,13 +64,16 @@ 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+
.. versionadded:: 2.0.0
7477
7578
Returns
7679
-------
@@ -119,6 +122,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
119122
>>> plt.show()
120123
121124
"""
125+
ifdevicenotin ["cpu",None]:
126+
raiseValueError(
127+
f"Unsupported device:{device}. Only\"cpu\" is allowed."
128+
)
129+
122130
num=operator.index(num)
123131
ifnum<0:
124132
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/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+
typedefenum {
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: 9 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-
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
"""
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,10 @@ 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+
.. versionadded:: 2.0.0
116123
117124
Returns
118125
-------

0 commit comments

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