8000 API: Add ``device`` and ``to_device`` to ``numpy.ndarray`` [Array API] by mtsokol · Pull Request #25233 · numpy/numpy · GitHub
[go: up one dir, main page]
Skip to content

API: Add device and to_device to numpy.ndarray [Array API] #25233

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
Jan 18, 2024
Merged

API: Add device and to_device to numpy.ndarray [Array API] #25233

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
26cd48f
API: Add device and to_device to ndarray
mtsokol Nov 23, 2023
6c51000
API: Add device keyword to fftfreq and rfftfreq
mtsokol Jan 10, 2024
897f9d0
Apply review comments
mtsokol Jan 17, 2024
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
Next Next commit
API: Add device and to_device to ndarray
  • Loading branch information
@mtsokol
mtsokol committed Jan 17, 2024
commit 26cd48fd452ea4f393aaa07e1d94983cbdcc9856
10 changes: 10 additions & 0 deletions doc/release/upcoming_changes/25233.new_feature.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
``ndarray.device`` and ``ndarray.to_device``
--------------------------------------------

``ndarray.device`` attribute and ``ndarray.to_device`` function were
added to `numpy.ndarray` class for Array API compatibility.

Additionally, ``device`` keyword-only arguments were added to:
`numpy.asarray`, `numpy.arange`, `numpy.empty`, `numpy.empty_like`,
`numpy.eye`, `numpy.full`, `numpy.full_like`, `numpy.linspace`,
`numpy.ones`, `numpy.ones_like`, `numpy.zeros`, and `numpy.zeros_like`.
6 changes: 0 additions & 6 deletions doc/source/reference/array_api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -501,12 +501,6 @@ Creation functions differences
* - ``copy`` keyword argument to ``asarray``
- **Compatible**
-
* - New ``device`` keyword argument to all array creation functions
(``asarray``, ``arange``, ``empty``, ``empty_like``, ``eye``, ``full``,
``full_like``, ``linspace``, ``ones``, ``ones_like``, ``zeros``, and
``zeros_like``).
- **Compatible**
- ``device`` would effectively do nothing, since NumPy is CPU only.

Elementwise functions differences
---------------------------------
Expand Down
5 changes: 5 additions & 0 deletions numpy/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2550,6 +2550,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):

def __array_namespace__(self, *, api_version: str = ...) -> Any: ...

def to_device(self, device: L["cpu"], /, *, stream: None | int | Any = ...) -> NDArray[Any]: ...

@property
def device(self) -> L["cpu"]: ...

def bitwise_count(
self,
out: None | NDArray[Any] = ...,
Expand Down
30 changes: 27 additions & 3 deletions numpy/_core/_add_newdocs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -912,7 +912,7 @@

add_newdoc('numpy._core.multiarray', 'asarray',
"""
asarray(a, dtype=None, order=None, *, like=None)
asarray(a, dtype=None, order=None, *, device=None, like=None)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here and below, why is device ahead of like? Since one cannot use it for anything useful, my tendency would be to put it at the end.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, why not just set the default to "cpu"? That seems clearer. (Still fine to accept None).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right now the function dispatch requires like to be the last argument - here's a check that raises RuntimeError if a function doesn't comply with it:

numpy/numpy/_core/overrides.py

Lines 149 to 156 in 448e4da

last_arg = co.co_argcount + co.co_kwonlyargcount - 1
last_arg = co.co_varnames[last_arg]
if last_arg != "like" or co.co_kwonlyargcount == 0:
raise RuntimeError(
"__array_function__ expects `like=` to be the last "
"argument and a keyword-only argument. "
f"{implementation} does not seem to comply.")

Can I keep device before like or should we update the __array_function__ protocol?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, why not just set the default to "cpu"? That seems clearer. (Still fine to accept None).

That's true that "cpu" is the only supported value, but I would prefer to stick to the Array API standard here. It defines the default value for device as None: https://data-apis.org/array-api/latest/API_specification/generated/array_api.asarray.html

We're adding these keyword arguments for Array API compatibility only, so I would prefer to fully conform with it.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it matters from the Array API perspective, it doesn't make NumPy non-compliant, it just makes it explicit about what that default is.

But, I'll give a very slight vote to just keep it None anyway. Just becuase I don't think the more explicit device="cpu" makes it clearer, since users should ignore it anyway (you only use it via xp = arr.__array_namespace__() as an Array API user, never directly as a NumPy user).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The order doesn't matter much in principle, since it's keyword-only and like is only for other non-numpy libraries as well. However, there's a bigger problem here:

>>> import numpy as np
>>> import dask.array as da
>>> d = da.ones((1, 2))
>>> np.asarray([1, 2], device='cpu')
array([1, 2])
>>> np.asarray([1, 2], device='cpu', like=d)
Traceback (most recent call last):
  Cell In[24], line 1
    np.asarray([1, 2], device='cpu', like=d)
  File ~/code/tmp/dask/dask/array/core.py:1760 in __array_function__
    return da_func(*args, **kwargs)
  File ~/code/tmp/dask/dask/array/core.py:4581 in asarray
    return from_array(a, getitem=getter_inline, **kwargs)
TypeError: from_array() got an unexpected keyword argument 'device'

The device keyword probably should not be added to the dispatchers for any of these creation functions. It looks like the __array_function__ protocol wasn't implemented in a forward-looking way - I haven't looked much deeper yet, but at first sight it may be the case that we can never add a new keyword to dispatched functions anymore.

In this case, device is here to support writing cross-library-compatible code via the array API standard. The dispatching protocols kinda were earlier attempts to do something similar. So probably they should not interact anyway?

Copy link
Member
@rgommers rgommers Jan 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah wait, it's maybe not that extreme - np.asarray([1, 2], like=d) still works, so new keywords are fine, as long as they are not used (then they don't get forwarded). Of course that's a bit funky, but there's no way to avoid it really.

>>> np.asarray([1, 2], device='cpu')
array([1, 2])
>>> np.asarray([1, 2], like=d)
dask.array<array, shape=(2,), dtype=int64, chunksize=(2,), chunktype=numpy.ndarray>
>>> np.asarray([1, 2], like=d, device='cpu')
...
TypeError: from_array() got an unexpected keyword argument 'device'

This would have been much more problematic if libraries like SciPy and scikit-learn supported non-numpy array inputs via __array_ufunc/function__. But since they squash such inputs anyway with np.asarray, there shouldn't be too much of an issue in practice.

And it's fairly easy to add support for new keywords in Dask & co. They just need a new 2.0-compatible release that adds such support, and then device/like can be used together.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@rgommers - I've gotten quite used to updating astropy's __array_function__ wrappers to remain compatible with numpy functions -- and we've got tests to check that our signatures remain consistent. Of course it doesn't happen that often. So, I don't think that's a worry, and overall I'd tend to not special-case any arguments (i.e., pass it on if present).

On the order: if like has to be last, then just keep it like that, sorry for not noticing myself.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to clarify: the check is only there to test habits, and probably mainly to make sure it is kwarg only.
It may have to be last in C (not sure), but it doesn't really matter anywhere here. But... it also doesn't matter a lot (especially since if someone cares, they can change it).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @mhvk, that is useful context. A test that flag mismatches sounds like a great idea. 🤞🏼 most implementers have that.


Convert the input to an array.

Expand All @@ -931,6 +931,14 @@
'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise
'K' (keep) preserve input order
Defaults to 'K'.
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}

.. versionadded:: 1.20.0
Expand Down Expand Up @@ -1184,7 +1192,7 @@

add_newdoc('numpy._core.multiarray', 'empty',
"""
empty(shape, dtype=float, order='C', *, like=None)
empty(shape, dtype=float, order='C', *, device=None, like=None)

Return a new array of given shape and type, without initializing entries.

Expand All @@ -1199,6 +1207,14 @@
Whether to store multi-dimensional data in row-major
(C-style) or column-major (Fortran-style) order in
memory.
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}

.. versionadded:: 1.20.0
Expand Down Expand Up @@ -1676,7 +1692,7 @@

add_newdoc('numpy._core.multiarray', 'arange',
"""
arange([start,] stop[, step,], dtype=None, *, like=None)
arange([start,] stop[, step,], dtype=None, *, device=None, like=None)

Return evenly spaced values within a given interval.

Expand Down Expand Up @@ -1717,6 +1733,14 @@
dtype : dtype, optional
The type of the output array. If `dtype` is not given, infer the data
type from the other input arguments.
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}

.. versionadded:: 1.20.0
Expand Down
18 changes: 15 additions & 3 deletions numpy/_core/function_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@


def _linspace_dispatcher(start, stop, num=None, endpoint=None, retstep=None,
dtype=None, axis=None):
dtype=None, axis=None, *, device=None):
return (start, stop)


@array_function_dispatch(_linspace_dispatcher)
def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
axis=0):
axis=0, *, device=None):
"""
Return evenly spaced numbers over a specified interval.

Expand Down Expand Up @@ -64,13 +64,20 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
array of integers.

.. versionadded:: 1.9.0

axis : int, optional
The axis in the result to store the samples. Relevant only if start
or stop are array-like. By default (0), the samples will be along a
new axis inserted at the beginning. Use -1 to get an axis at the end.

.. versionadded:: 1.16.0
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.

.. versionadded:: 2.0.0

Returns
-------
Expand Down Expand Up @@ -119,6 +126,11 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
>>> plt.show()

"""
if device not in ["cpu", None]:
raise ValueError(
f"Unsupported device: {device}. Only \"cpu\" is allowed."
)

num = operator.index(num)
if num < 0:
raise ValueError(
Expand Down
16 changes: 16 additions & 0 deletions numpy/_core/function_base.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ def linspace(
retstep: L[False] = ...,
dtype: None = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> NDArray[floating[Any]]: ...
@overload
def linspace(
Expand All @@ -38,6 +40,8 @@ def linspace(
retstep: L[False] = ...,
dtype: None = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> NDArray[complexfloating[Any, Any]]: ...
@overload
def linspace(
Expand All @@ -48,6 +52,8 @@ def linspace(
retstep: L[False] = ...,
dtype: _DTypeLike[_SCT] = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> NDArray[_SCT]: ...
@overload
def linspace(
Expand All @@ -58,6 +64,8 @@ def linspace(
retstep: L[False] = ...,
dtype: DTypeLike = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> NDArray[Any]: ...
@overload
def linspace(
Expand All @@ -68,6 +76,8 @@ def linspace(
retstep: L[True] = ...,
dtype: None = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> tuple[NDArray[floating[Any]], floating[Any]]: ...
@overload
def linspace(
Expand All @@ -78,6 +88,8 @@ def linspace(
retstep: L[True] = ...,
dtype: None = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> tuple[NDArray[complexfloating[Any, Any]], complexfloating[Any, Any]]: ...
@overload
def linspace(
Expand All @@ -88,6 +100,8 @@ def linspace(
retstep: L[True] = ...,
dtype: _DTypeLike[_SCT] = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> tuple[NDArray[_SCT], _SCT]: ...
@overload
def linspace(
Expand All @@ -98,6 +112,8 @@ def linspace(
retstep: L[True] = ...,
dtype: DTypeLike = ...,
axis: SupportsIndex = ...,
*,
device: None | L["cpu"] = ...,
) -> tuple[NDArray[Any], Any]: ...

@overload
Expand Down
5 changes: 5 additions & 0 deletions numpy/_core/include/numpy/ndarraytypes.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,11 @@ typedef enum {
NPY_BUSDAY_RAISE
} NPY_BUSDAY_ROLL;

/* Device enum for Array API compatibility */
typedef enum {
NPY_DEVICE_CPU = 0,
} NPY_DEVICE;

/************************************************************
* NumPy Auxiliary Data for inner loops, sort functions, etc.
************************************************************/
Expand Down
15 changes: 13 additions & 2 deletions numpy/_core/multiarray.py
View file
Open in desktop
AD5A
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,12 @@


@array_function_from_c_func_and_dispatcher(_multiarray_umath.empty_like)
def empty_like(prototype, dtype=None, order=None, subok=None, shape=None):
def empty_like(
prototype, dtype=None, order=None, subok=None, shape=None, *, device=None
):
"""
empty_like(prototype, dtype=None, order='K', subok=True, shape=None)
empty_like(prototype, dtype=None, order='K', subok=True, shape=None,
shape=None, *, device=None)

Return a new array with the same shape and type as a given array.

Expand Down Expand Up @@ -113,6 +116,14 @@ def empty_like(prototype, dtype=None, order=None, subok=None, shape=None):
order='C' is implied.

.. versionadded:: 1.17.0
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.

.. versionadded:: 2.0.0

Returns
-------
Expand Down
Loading
0