diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py
index 9f8be2479..9f8f3ca4d 100644
--- a/src/array_api_stubs/_draft/array_object.py
+++ b/src/array_api_stubs/_draft/array_object.py
@@ -293,6 +293,8 @@ def __dlpack__(
*,
stream: Optional[Union[int, Any]] = None,
max_version: Optional[tuple[int, int]] = None,
+ dl_device: Optional[tuple[Enum, int]] = None,
+ copy: Optional[bool] = None,
) -> PyCapsule:
"""
Exports the array for consumption by :func:`~array_api.from_dlpack` as a DLPack capsule.
@@ -324,6 +326,12 @@ def __dlpack__(
- ``> 2``: stream number represented as a Python integer.
- Using ``1`` and ``2`` is not supported.
+ .. note::
+ When ``dl_device`` is provided explicitly, ``stream`` must be a valid
+ construct for the specified device type. In particular, when ``kDLCPU``
+ is in use, ``stream`` must be ``None`` and a synchronization must be
+ performed to ensure data safety.
+
.. admonition:: Tip
:class: important
@@ -333,12 +341,40 @@ def __dlpack__(
not want to think about stream handling at all, potentially at the
cost of more synchronizations than necessary.
max_version: Optional[tuple[int, int]]
- The maximum DLPack version that the *consumer* (i.e., the caller of
+ the maximum DLPack version that the *consumer* (i.e., the caller of
``__dlpack__``) supports, in the form of a 2-tuple ``(major, minor)``.
This method may return a capsule of version ``max_version`` (recommended
if it does support that), or of a different version.
This means the consumer must verify the version even when
`max_version` is passed.
+ dl_device: Optional[tuple[enum.Enum, int]]
+ the DLPack device type. Default is ``None``, meaning the exported capsule
+ should be on the same device as ``self`` is. When specified, the format
+ must be a 2-tuple, following that of the return value of :meth:`array.__dlpack_device__`.
+ If the device type cannot be handled by the producer, this function must
+ raise ``BufferError``.
+
+ The v2023.12 standard only mandates that a compliant library should offer a way for
+ ``__dlpack__`` to return a capsule referencing an array whose underlying memory is
+ accessible to the Python interpreter (represented by the ``kDLCPU`` enumerator in DLPack).
+ If a copy must be made to enable this support but ``copy`` is set to ``False``, the
+ function must raise ``ValueError``.
+
+ Other device kinds will be considered for standardization in a future version of this
+ API standard.
+ copy: Optional[bool]
+ boolean indicating whether or not to copy the input. If ``True``, the
+ function must always copy (performed by the producer). If ``False``, the
+ function must never copy, and raise a ``BufferError`` in case a copy is
+ deemed necessary (e.g. if a cross-device data movement is requested, and
+ it is not possible without a copy). If ``None``, the function must reuse
+ the existing memory buffer if possible and copy otherwise. Default: ``None``.
+
+ When a copy happens, the ``DLPACK_FLAG_BITMASK_IS_COPIED`` flag must be set.
+
+ .. note::
+ If a copy happens, and if the consumer-provided ``stream`` and ``dl_device``
+ can be understood by the producer, the copy must be performed over ``stream``.
Returns
-------
@@ -394,22 +430,25 @@ def __dlpack__(
# here to tell users that the consumer's max_version is too
# old to allow the data exchange to happen.
- And this logic for the consumer in ``from_dlpack``:
+ And this logic for the consumer in :func:`~array_api.from_dlpack`:
.. code:: python
try:
- x.__dlpack__(max_version=(1, 0))
+ x.__dlpack__(max_version=(1, 0), ...)
# if it succeeds, store info from the capsule named "dltensor_versioned",
# and need to set the name to "used_dltensor_versioned" when we're done
except TypeError:
- x.__dlpack__()
+ x.__dlpack__(...)
+
+ This logic is also applicable to handling of the new ``dl_device`` and ``copy``
+ keywords.
.. versionchanged:: 2022.12
Added BufferError.
.. versionchanged:: 2023.12
- Added the ``max_version`` keyword.
+ Added the ``max_version``, ``dl_device``, and ``copy`` keywords.
"""
def __dlpack_device__(self: array, /) -> Tuple[Enum, int]:
@@ -436,6 +475,8 @@ def __dlpack_device__(self: array, /) -> Tuple[Enum, int]:
METAL = 8
VPI = 9
ROCM = 10
+ CUDA_MANAGED = 13
+ ONE_API = 14
"""
def __eq__(self: array, other: Union[int, float, bool, array], /) -> array:
diff --git a/src/array_api_stubs/_draft/creation_functions.py b/src/array_api_stubs/_draft/creation_functions.py
index d0f967717..dec09db0c 100644
--- a/src/array_api_stubs/_draft/creation_functions.py
+++ b/src/array_api_stubs/_draft/creation_functions.py
@@ -214,7 +214,13 @@ def eye(
"""
-def from_dlpack(x: object, /) -> array:
+def from_dlpack(
+ x: object,
+ /,
+ *,
+ device: Optional[device] = None,
+ copy: Optional[bool] = None,
+) -> array:
"""
Returns a new array containing the data from another (array) object with a ``__dlpack__`` method.
@@ -222,11 +228,22 @@ def from_dlpack(x: object, /) -> array:
----------
x: object
input (array) object.
+ device: Optional[device]
+ device on which to place the created array. If ``device`` is ``None`` and ``x`` supports DLPack, the output array must be on the same device as ``x``. Default: ``None``.
+
+ The v2023.12 standard only mandates that a compliant library should offer a way for ``from_dlpack`` to return an array
+ whose underlying memory is accessible to the Python interpreter, when the corresponding ``device`` is provided. If the
+ array library does not support such cases at all, the function must raise ``BufferError``. If a copy must be made to
+ enable this support but ``copy`` is set to ``False``, the function must raise ``ValueError``.
+
+ Other device kinds will be considered for standardization in a future version of this API standard.
+ copy: Optional[bool]
+ boolean indicating whether or not to copy the input. If ``True``, the function must always copy. If ``False``, the function must never copy, and raise ``BufferError`` in case a copy is deemed necessary (e.g. if a cross-device data movement is requested, and it is not possible without a copy). If ``None``, the function must reuse the existing memory buffer if possible and copy otherwise. Default: ``None``.
Returns
-------
out: array
- an array containing the data in `x`.
+ an array containing the data in ``x``.
.. admonition:: Note
:class: note
@@ -238,7 +255,7 @@ def from_dlpack(x: object, /) -> array:
BufferError
The ``__dlpack__`` and ``__dlpack_device__`` methods on the input array
may raise ``BufferError`` when the data cannot be exported as DLPack
- (e.g., incompatible dtype or strides). It may also raise other errors
+ (e.g., incompatible dtype, strides, or device). It may also raise other errors
when export fails for other reasons (e.g., not enough memory available
to materialize the data). ``from_dlpack`` must propagate such
exceptions.
@@ -246,11 +263,34 @@ def from_dlpack(x: object, /) -> array:
If the ``__dlpack__`` and ``__dlpack_device__`` methods are not present
on the input array. This may happen for libraries that are never able
to export their data with DLPack.
+ ValueError
+ If data exchange is possible via an explicit copy but ``copy`` is set to ``False``.
Notes
-----
See :meth:`array.__dlpack__` for implementation suggestions for `from_dlpack` in
order to handle DLPack versioning correctly.
+
+ A way to move data from two array libraries to the same device (assumed supported by both libraries) in
+ a library-agnostic fashion is illustrated below:
+
+ .. code:: python
+
+ def func(x, y):
+ xp_x = x.__array_namespace__()
+ xp_y = y.__array_namespace__()
+
+ # Other functions than `from_dlpack` only work if both arrays are from the same library. So if
+ # `y` is from a different one than `x`, let's convert `y` into an array of the same type as `x`:
+ if not xp_x == xp_y:
+ y = xp_x.from_dlpack(y, copy=True, device=x.device)
+
+ # From now on use `xp_x.xxxxx` functions, as both arrays are from the library `xp_x`
+ ...
+
+
+ .. versionchanged:: 2023.12
+ Added device and copy support.
"""