Skip to content

MAINT: __array_namespace_info__ docstrings tweaks #300

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 15, 2025
Merged

MAINT: __array_namespace_info__ docstrings tweaks #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7e4d303
MAINT: clarify `default_device` output
crusaderky Apr 3, 2025
04e2268
Remove outdated comment on 'max rank'
crusaderky Apr 11, 2025
4172d85
Update docs
crusaderky Apr 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion array_api_compat/common/_aliases.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

# These functions are modified from the NumPy versions.

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

def arange(
start: Union[int, float],
Expand Down
20 changes: 15 additions & 5 deletions array_api_compat/cupy/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
complex128,
)


class __array_namespace_info__:
"""
Get the array API inspection namespace for CuPy.
Expand All @@ -49,7 +50,7 @@ class __array_namespace_info__:

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_dtypes()
{'real floating': cupy.float64,
'complex floating': cupy.complex128,
Expand Down Expand Up @@ -94,13 +95,13 @@ def capabilities(self):
>>> info = xp.__array_namespace_info__()
>>> info.capabilities()
{'boolean indexing': True,
'data-dependent shapes': True}
'data-dependent shapes': True,
'max dimensions': 64}

"""
return {
"boolean indexing": True,
"data-dependent shapes": True,
# 'max rank' will be part of the 2024.12 standard
"max dimensions": 64,
}

Expand All @@ -117,7 +118,7 @@ def default_device(self):

Returns
-------
device : str
device : Device
The default device used for new CuPy arrays.

Examples
Expand All @@ -126,6 +127,15 @@ def default_device(self):
>>> info.default_device()
Device(0)

Notes
-----
This method returns the static default device when CuPy is initialized.
However, the *current* device used by creation functions (``empty`` etc.)
can be changed globally or with a context manager.

See Also
--------
https://github.com/data-apis/array-api/issues/835
"""
return cuda.Device(0)

Expand Down Expand Up @@ -312,7 +322,7 @@ def devices(self):

Returns
-------
devices : list of str
devices : list[Device]
The devices supported by CuPy.

See Also
Expand Down
19 changes: 10 additions & 9 deletions array_api_compat/dask/array/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ class __array_namespace_info__:

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_dtypes()
{'real floating': dask.float64,
'complex floating': dask.complex128,
Expand Down Expand Up @@ -103,10 +103,11 @@ def capabilities(self):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.capabilities()
{'boolean indexing': True,
'data-dependent shapes': True}
'data-dependent shapes': True,
'max dimensions': 64}

"""
return {
Expand All @@ -130,12 +131,12 @@ def default_device(self):

Returns
-------
device : str
device : Device
The default device used for new Dask arrays.

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_device()
'cpu'

Expand Down Expand Up @@ -173,7 +174,7 @@ def default_dtypes(self, *, device=None):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_dtypes()
{'real floating': dask.float64,
'complex floating': dask.complex128,
Expand Down Expand Up @@ -239,7 +240,7 @@ def dtypes(self, *, device=None, kind=None):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.dtypes(kind='signed integer')
{'int8': dask.int8,
'int16': dask.int16,
Expand Down Expand Up @@ -335,7 +336,7 @@ def devices(self):

Returns
-------
devices : list of str
devices : list[Device]
The devices supported by Dask.

See Also
Expand All @@ -347,7 +348,7 @@ def devices(self):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.devices()
['cpu', DASK_DEVICE]

Expand Down
8 changes: 4 additions & 4 deletions array_api_compat/numpy/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,13 @@ def capabilities(self):
>>> info = np.__array_namespace_info__()
>>> info.capabilities()
{'boolean indexing': True,
'data-dependent shapes': True}
'data-dependent shapes': True,
'max dimensions': 64}

"""
return {
"boolean indexing": True,
"data-dependent shapes": True,
# 'max rank' will be part of the 2024.12 standard
"max dimensions": 64,
}

Expand All @@ -119,7 +119,7 @@ def default_device(self):

Returns
-------
device : str
device : Device
The default device used for new NumPy arrays.

Examples
Expand Down Expand Up @@ -326,7 +326,7 @@ def devices(self):

Returns
-------
devices : list of str
devices : list[Device]
The devices supported by NumPy.

See Also
Expand Down
41 changes: 26 additions & 15 deletions array_api_compat/torch/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ class __array_namespace_info__:

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_dtypes()
{'real floating': numpy.float64,
'complex floating': numpy.complex128,
Expand Down Expand Up @@ -76,16 +76,16 @@ def capabilities(self):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.capabilities()
{'boolean indexing': True,
'data-dependent shapes': True}
'data-dependent shapes': True,
'max dimensions': 64}

"""
return {
"boolean indexing": True,
"data-dependent shapes": True,
# 'max rank' will be part of the 2024.12 standard
"max dimensions": 64,
}

Expand All @@ -102,15 +102,24 @@ def default_device(self):

Returns
-------
device : str
device : Device
The default device used for new PyTorch arrays.

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_device()
'cpu'
device(type='cpu')

Notes
-----
This method returns the static default device when PyTorch is initialized.
However, the *current* device used by creation functions (``empty`` etc.)
can be changed at runtime.

See Also
--------
https://github.com/data-apis/array-api/issues/835
"""
return torch.device("cpu")

Expand All @@ -120,9 +129,9 @@ def default_dtypes(self, *, device=None):

Parameters
----------
device : str, optional
The device to get the default data types for. For PyTorch, only
``'cpu'`` is allowed.
device : Device, optional
The device to get the default data types for.
Unused for PyTorch, as all devices use the same default dtypes.

Returns
-------
Expand All @@ -139,7 +148,7 @@ def default_dtypes(self, *, device=None):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.default_dtypes()
{'real floating': torch.float32,
'complex floating': torch.complex64,
Expand Down Expand Up @@ -250,8 +259,9 @@ def dtypes(self, *, device=None, kind=None):

Parameters
----------
device : str, optional
device : Device, optional
The device to get the data types for.
Unused for PyTorch, as all devices use the same dtypes.
kind : str or tuple of str, optional
The kind of data types to return. If ``None``, all data types are
returned. If a string, only data types of that kind are returned.
Expand Down Expand Up @@ -287,7 +297,7 @@ def dtypes(self, *, device=None, kind=None):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.dtypes(kind='signed integer')
{'int8': numpy.int8,
'int16': numpy.int16,
Expand All @@ -310,7 +320,7 @@ def devices(self):

Returns
-------
devices : list of str
devices : list[Device]
The devices supported by PyTorch.

See Also
Expand All @@ -322,7 +332,7 @@ def devices(self):

Examples
--------
>>> info = np.__array_namespace_info__()
>>> info = xp.__array_namespace_info__()
>>> info.devices()
[device(type='cpu'), device(type='mps', index=0), device(type='meta')]

Expand All @@ -333,6 +343,7 @@ def devices(self):
# device:
try:
torch.device('notadevice')
raise AssertionError("unreachable") # pragma: nocover
except RuntimeError as e:
# The error message is something like:
# "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"
Expand Down
Loading