10000 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
Prev Previous commit
Apply review comments
  • Loading branch information
@mtsokol
mtsokol committed Jan 17, 2024
commit 897f9d096328142a2f9d8bf2954c19d35ee64d46
4 changes: 3 additions & 1 deletion doc/release/upcoming_changes/25233.new_feature.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
``ndarray.device`` and ``ndarray.to_device``
--------------------------------------------

``ndarray.device`` attribute and ``ndarray.to_device`` function were
``ndarray.device`` attribute and ``ndarray.to_device`` method 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`.

For all these new arguments, only ``device="cpu"`` is supported.
19 changes: 5 additions & 14 deletions numpy/_core/_add_newdocs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -933,10 +933,7 @@
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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand All @@ -946,8 +943,8 @@
Returns
-------
out : ndarray
Array interpretation of `a`. No copy is performed if the input
is already an ndarray with matching dtype and order. If `a` is a
Array interpretation of ``a``. No copy is performed if the input
is already an ndarray with matching dtype and order. If ``a`` is a
subclass of ndarray, a base class ndarray is returned.

See Also
Expand Down Expand Up @@ -1209,10 +1206,7 @@
memory.
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand Down Expand Up @@ -1735,10 +1729,7 @@
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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand Down
14 changes: 4 additions & 10 deletions numpy/_core/function_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,10 +72,7 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
.. 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand Down Expand Up @@ -126,11 +123,6 @@ 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 All @@ -152,7 +144,9 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
integer_dtype = _nx.issubdtype(dtype, _nx.integer)

delta = stop - start
y = _nx.arange(0, num, dtype=dt).reshape((-1,) + (1,) * ndim(delta))
y = _nx.arange(
0, num, dtype=dt, device=device
).reshape((-1,) + (1,) * ndim(delta))
# In-place multiplication y *= delta/div is faster, but prevents
# the multiplicant from overriding what class is produced, and thus
# prevents, e.g. use of Quantities, see gh-7142. Hence, we multiply
Expand Down
5 changes: 1 addition & 4 deletions numpy/_core/multiarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,10 +118,7 @@ def empty_like(
.. 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand Down
55 changes: 15 additions & 40 deletions numpy/_core/numeric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,10 +96,7 @@ def zeros_like(
.. 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand Down Expand Up @@ -161,10 +158,7 @@ def ones(shape, dtype=None, order='C', *, device=None, like=None):
memory.
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand Down Expand Up @@ -201,15 +195,12 @@ def ones(shape, dtype=None, order='C', *, device=None, like=None):
[1., 1.]])

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

if like is not None:
return _ones_with_like(like, shape, dtype=dtype, order=order)
return _ones_with_like(
like, shape, dtype=dtype, order=order, device=device
)

a = empty(shape, dtype, order)
a = empty(shape, dtype, order, device=device)
multiarray.copyto(a, 1, casting='unsafe')
return a

Expand Down Expand Up @@ -258,10 +249,7 @@ def ones_like(
.. 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand Down Expand Up @@ -328,10 +316,7 @@ def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None):
(row- or column-wise) 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand Down Expand Up @@ -364,19 +349,15 @@ def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None):
[1, 2]])

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

if like is not None:
return _full_with_like(
like, shape, fill_value, dtype=dtype, order=order)
like, shape, fill_value, dtype=dtype, order=order, device=device
)

if dtype is None:
fill_value = asarray(fill_value)
dtype = fill_value.dtype
a = empty(shape, dtype, order)
a = empty(shape, dtype, order, device=device)
multiarray.copyto(a, fill_value, casting='unsafe')
return a

Expand Down Expand Up @@ -425,10 +406,7 @@ def full_like(
.. 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.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand Down Expand Up @@ -467,12 +445,9 @@ def full_like(
[[ 0, 0, 255],
[ 0, 0, 255]]])
"""
if device not in ["cpu", None]:
raise ValueError(
f"Unsupported device: {device}. Only \"cpu\" is allowed."
)

res = empty_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
res = empty_like(
a, dtype=dtype, order=order, subok=subok, shape=shape, device=device
)
multiarray.copyto(res, fill_value, casting='unsafe')
return res

Expand Down
5 changes: 3 additions & 2 deletions numpy/_core/src/multiarray/conversion_utils.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1370,7 +1370,8 @@ PyArray_DeviceConverterOptional(PyObject *object, NPY_DEVICE *device)
return NPY_SUCCEED;
}

PyErr_SetString(PyExc_ValueError,
"Device not understood. Only \"cpu\" is allowed.");
PyErr_Format(PyExc_ValueError,
"Device not understood. Only \"cpu\" is allowed, "
"but received: %S", object);
return NPY_FAIL;
}
2 changes: 2 additions & 0 deletions numpy/_core/src/multiarray/multiarraymodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2181,13 +2181,15 @@ array_zeros(PyObject *NPY_UNUSED(ignored),
NPY_ORDER order = NPY_CORDER;
npy_bool is_f_order = NPY_FALSE;
PyArrayObject *ret = NULL;
NPY_DEVICE device = NPY_DEVICE_CPU;
PyObject *like = Py_None;
NPY_PREPARE_ARGPARSER;

if (npy_parse_arguments("zeros", args, len_args, kwnames,
"shape", &PyArray_IntpConverter, &shape,
"|dtype", &PyArray_DTypeOrDescrConverterOptional, &dt_info,
"|order", &PyArray_OrderConverter, &order,
"$device", &PyArray_DeviceConverterOptional, &device,
"$like", NULL, &like,
NULL, NULL, NULL) < 0) {
goto finish;
Expand Down
24 changes: 6 additions & 18 deletions numpy/fft/_helper.py
View file
Open in desktop
10000
Original file line number Diff line number Diff line change
Expand Up @@ -142,9 +142,7 @@ def fftfreq(n, d=1.0, device=None):
Sample spacing (inverse of the sampling rate). Defaults to 1.
device : str, optional
The device on which to place the created array. Default: ``None``.

.. note::
Only the ``"cpu"`` device is supported by NumPy.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand All @@ -164,18 +162,14 @@ def fftfreq(n, d=1.0, device=None):
array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25])

"""
if device not in ["cpu", None]:
raise ValueError(
f"Unsupported device: {device}. Only \"cpu\" is allowed."
)
if not isinstance(n, integer_types):
raise ValueError("n should be an integer")
val = 1.0 / (n * d)
results = empty(n, int)
results = empty(n, int, device=device)
N = (n-1)//2 + 1
p1 = arange(0, N, dtype=int)
p1 = arange(0, N, dtype=int, device=device)
results[:N] = p1
p2 = arange(-(n//2), 0, dtype=int)
p2 = arange(-(n//2), 0, dtype=int, device=device)
results[N:] = p2
return results * val

Expand Down Expand Up @@ -206,9 +200,7 @@ def rfftfreq(n, d=1.0, device=None):
Sample spacing (inverse of the sampling rate). Defaults to 1.
device : str, optional
The device on which to place the created array. Default: ``None``.

.. note::
Only the ``"cpu"`` device is supported by NumPy.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0

Expand All @@ -231,13 +223,9 @@ def rfftfreq(n, d=1.0, device=None):
array([ 0., 10., 20., 30., 40., 50.])

"""
if device not in ["cpu", None]:
raise ValueError(
f"Unsupported device: {device}. Only \"cpu\" is allowed."
)
if not isinstance(n, integer_types):
raise ValueError("n should be an integer")
val = 1.0/(n*d)
N = n//2 + 1
results = arange(0, N, dtype=int)
results = arange(0, N, dtype=int, device=device)
return results * val
16 changes: 5 additions & 11 deletions numpy/lib/_twodim_base_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,7 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None):
.. versionadded:: 1.14.0
device : str, optional
The device on which to place the created array. Default: None.

.. note::

Only the ``"cpu"`` device is supported by NumPy.
For Array-API interoperability only, so must be ``"cpu"`` if passed.

.. versionadded:: 2.0.0
${ARRAY_FUNCTION_LIKE}
Expand Down Expand Up @@ -212,16 +209,13 @@ def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None):
[0., 0., 0.]])

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

if like is not None:
return _eye_with_like(like, N, M=M, k=k, dtype=dtype, order=order)
return _eye_with_like(
like, N, M=M, k=k, dtype=dtype, order=order, device=device
)
if M is None:
M = N
m = zeros((N, M), dtype=dtype, order=order)
m = zeros((N, M), dtype=dtype, order=order, device=device)
if k >= M:
return m
# Ensure M and k are integers, so we don't get any surprise casting
Expand Down
0