Skip to content

Commit d743dc1

Browse files
crusaderkycrusaderky
authored
MAINT: __array_namespace_info__ docstrings tweaks (#300)
1 parent bff3bf4 commit d743dc1

File tree

5 files changed

+56
-34
lines changed

5 files changed

+56
-34
lines changed

‎array_api_compat/common/_aliases.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
# These functions are modified from the NumPy versions.
2020

21-
# Creation functions add the device keyword (which does nothing for NumPy)
21+
# Creation functions add the device keyword (which does nothing for NumPy and Dask)
2222

2323
def arange(
2424
start: Union[int, float],

‎array_api_compat/cupy/_info.py

+15-5
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
complex128,
2727
)
2828

29+
2930
class __array_namespace_info__:
3031
"""
3132
Get the array API inspection namespace for CuPy.
@@ -49,7 +50,7 @@ class __array_namespace_info__:
4950
5051
Examples
5152
--------
52-
>>> info = np.__array_namespace_info__()
53+
>>> info = xp.__array_namespace_info__()
5354
>>> info.default_dtypes()
5455
{'real floating': cupy.float64,
5556
'complex floating': cupy.complex128,
@@ -94,13 +95,13 @@ def capabilities(self):
9495
>>> info = xp.__array_namespace_info__()
9596
>>> info.capabilities()
9697
{'boolean indexing': True,
97-
'data-dependent shapes': True}
98+
'data-dependent shapes': True,
99+
'max dimensions': 64}
98100
99101
"""
100102
return {
101103
"boolean indexing": True,
102104
"data-dependent shapes": True,
103-
# 'max rank' will be part of the 2024.12 standard
104105
"max dimensions": 64,
105106
}
106107

@@ -117,7 +118,7 @@ def default_device(self):
117118
118119
Returns
119120
-------
120-
device : str
121+
device : Device
121122
The default device used for new CuPy arrays.
122123
123124
Examples
@@ -126,6 +127,15 @@ def default_device(self):
126127
>>> info.default_device()
127128
Device(0)
128129
130+
Notes
131+
-----
132+
This method returns the static default device when CuPy is initialized.
133+
However, the *current* device used by creation functions (``empty`` etc.)
134+
can be changed globally or with a context manager.
135+
136+
See Also
137+
--------
138+
https://github.com/data-apis/array-api/issues/835
129139
"""
130140
return cuda.Device(0)
131141

@@ -312,7 +322,7 @@ def devices(self):
312322
313323
Returns
314324
-------
315-
devices : list of str
325+
devices : list[Device]
316326
The devices supported by CuPy.
317327
318328
See Also

‎array_api_compat/dask/array/_info.py

+10-9
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ class __array_namespace_info__:
5050
5151
Examples
5252
--------
53-
>>> info = np.__array_namespace_info__()
53+
>>> info = xp.__array_namespace_info__()
5454
>>> info.default_dtypes()
5555
{'real floating': dask.float64,
5656
'complex floating': dask.complex128,
@@ -103,10 +103,11 @@ def capabilities(self):
103103
104104
Examples
105105
--------
106-
>>> info = np.__array_namespace_info__()
106+
>>> info = xp.__array_namespace_info__()
107107
>>> info.capabilities()
108108
{'boolean indexing': True,
109-
'data-dependent shapes': True}
109+
'data-dependent shapes': True,
110+
'max dimensions': 64}
110111
111112
"""
112113
return {
@@ -130,12 +131,12 @@ def default_device(self):
130131
131132
Returns
132133
-------
133-
device : str
134+
device : Device
134135
The default device used for new Dask arrays.
135136
136137
Examples
137138
--------
138-
>>> info = np.__array_namespace_info__()
139+
>>> info = xp.__array_namespace_info__()
139140
>>> info.default_device()
140141
'cpu'
141142
@@ -173,7 +174,7 @@ def default_dtypes(self, *, device=None):
173174
174175
Examples
175176
--------
176-
>>> info = np.__array_namespace_info__()
177+
>>> info = xp.__array_namespace_info__()
177178
>>> info.default_dtypes()
178179
{'real floating': dask.float64,
179180
'complex floating': dask.complex128,
@@ -239,7 +240,7 @@ def dtypes(self, *, device=None, kind=None):
239240
240241
Examples
241242
--------
242-
>>> info = np.__array_namespace_info__()
243+
>>> info = xp.__array_namespace_info__()
243244
>>> info.dtypes(kind='signed integer')
244245
{'int8': dask.int8,
245246
'int16': dask.int16,
@@ -335,7 +336,7 @@ def devices(self):
335336
336337
Returns
337338
-------
338-
devices : list of str
339+
devices : list[Device]
339340
The devices supported by Dask.
340341
341342
See Also
@@ -347,7 +348,7 @@ def devices(self):
347348
348349
Examples
349350
--------
350-
>>> info = np.__array_namespace_info__()
351+
>>> info = xp.__array_namespace_info__()
351352
>>> info.devices()
352353
['cpu', DASK_DEVICE]
353354

‎array_api_compat/numpy/_info.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -94,13 +94,13 @@ def capabilities(self):
9494
>>> info = np.__array_namespace_info__()
9595
>>> info.capabilities()
9696
{'boolean indexing': True,
97-
'data-dependent shapes': True}
97+
'data-dependent shapes': True,
98+
'max dimensions': 64}
9899
99100
"""
100101
return {
101102
"boolean indexing": True,
102103
"data-dependent shapes": True,
103-
# 'max rank' will be part of the 2024.12 standard
104104
"max dimensions": 64,
105105
}
106106

@@ -119,7 +119,7 @@ def default_device(self):
119119
120120
Returns
121121
-------
122-
device : str
122+
device : Device
123123
The default device used for new NumPy arrays.
124124
125125
Examples
@@ -326,7 +326,7 @@ def devices(self):
326326
327327
Returns
328328
-------
329-
devices : list of str
329+
devices : list[Device]
330330
The devices supported by NumPy.
331331
332332
See Also

‎array_api_compat/torch/_info.py

+26-15
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ class __array_namespace_info__:
3434
3535
Examples
3636
--------
37-
>>> info = np.__array_namespace_info__()
37+
>>> info = xp.__array_namespace_info__()
3838
>>> info.default_dtypes()
3939
{'real floating': numpy.float64,
4040
'complex floating': numpy.complex128,
@@ -76,16 +76,16 @@ def capabilities(self):
7676
7777
Examples
7878
--------
79-
>>> info = np.__array_namespace_info__()
79+
>>> info = xp.__array_namespace_info__()
8080
>>> info.capabilities()
8181
{'boolean indexing': True,
82-
'data-dependent shapes': True}
82+
'data-dependent shapes': True,
83+
'max dimensions': 64}
8384
8485
"""
8586
return {
8687
"boolean indexing": True,
8788
"data-dependent shapes": True,
88-
# 'max rank' will be part of the 2024.12 standard
8989
"max dimensions": 64,
9090
}
9191

@@ -102,15 +102,24 @@ def default_device(self):
102102
103103
Returns
104104
-------
105-
device : str
105+
device : Device
106106
The default device used for new PyTorch arrays.
107107
108108
Examples
109109
--------
110-
>>> info = np.__array_namespace_info__()
110+
>>> info = xp.__array_namespace_info__()
111111
>>> info.default_device()
112-
'cpu'
112+
device(type='cpu')
113113
114+
Notes
115+
-----
116+
This method returns the static default device when PyTorch is initialized.
117+
However, the *current* device used by creation functions (``empty`` etc.)
118+
can be changed at runtime.
119+
120+
See Also
121+
--------
122+
https://github.com/data-apis/array-api/issues/835
114123
"""
115124
return torch.device("cpu")
116125

@@ -120,9 +129,9 @@ def default_dtypes(self, *, device=None):
120129
121130
Parameters
122131
----------
123-
device : str, optional
124-
The device to get the default data types for. For PyTorch, only
125-
``'cpu'`` is allowed.
132+
device : Device, optional
133+
The device to get the default data types for.
134+
Unused for PyTorch, as all devices use the same default dtypes.
126135
127136
Returns
128137
-------
@@ -139,7 +148,7 @@ def default_dtypes(self, *, device=None):
139148
140149
Examples
141150
--------
142-
>>> info = np.__array_namespace_info__()
151+
>>> info = xp.__array_namespace_info__()
143152
>>> info.default_dtypes()
144153
{'real floating': torch.float32,
145154
'complex floating': torch.complex64,
@@ -250,8 +259,9 @@ def dtypes(self, *, device=None, kind=None):
250259
251260
Parameters
252261
----------
253-
device : str, optional
262+
device : Device, optional
254263
The device to get the data types for.
264+
Unused for PyTorch, as all devices use the same dtypes.
255265
kind : str or tuple of str, optional
256266
The kind of data types to return. If ``None``, all data types are
257267
returned. If a string, only data types of that kind are returned.
@@ -287,7 +297,7 @@ def dtypes(self, *, device=None, kind=None):
287297
288298
Examples
289299
--------
290-
>>> info = np.__array_namespace_info__()
300+
>>> info = xp.__array_namespace_info__()
291301
>>> info.dtypes(kind='signed integer')
292302
{'int8': numpy.int8,
293303
'int16': numpy.int16,
@@ -310,7 +320,7 @@ def devices(self):
310320
311321
Returns
312322
-------
313-
devices : list of str
323+
devices : list[Device]
314324
The devices supported by PyTorch.
315325
316326
See Also
@@ -322,7 +332,7 @@ def devices(self):
322332
323333
Examples
324334
--------
325-
>>> info = np.__array_namespace_info__()
335+
>>> info = xp.__array_namespace_info__()
326336
>>> info.devices()
327337
[device(type='cpu'), device(type='mps', index=0), device(type='meta')]
328338
@@ -333,6 +343,7 @@ def devices(self):
333343
# device:
334344
try:
335345
torch.device('notadevice')
346+
raise AssertionError("unreachable") # pragma: nocover
336347
except RuntimeError as e:
337348
# The error message is something like:
338349
# "Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: notadevice"

0 commit comments

Comments
 (0)