Skip to content
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

MAINT: clarify default_device output #300

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

MAINT: clarify default_device output #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 12 additions & 2 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 Down Expand Up @@ -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
4 changes: 2 additions & 2 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 @@ -130,7 +130,7 @@ def default_device(self):

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

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

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

See Also
Expand Down
4 changes: 2 additions & 2 deletions array_api_compat/numpy/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
25 changes: 18 additions & 7 deletions array_api_compat/torch/_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.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 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 @@ -310,7 +320,7 @@ def devices(self):

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

See Also
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