Update core.py

pydata · Aug 20, 2024 · d0dd5f4 · d0dd5f4
1 parent 3f2e2f9
commit d0dd5f4
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py
@@ -50,6 +50,7 @@
 )
 
 if TYPE_CHECKING:
+    from enum import IntEnum
     from numpy.typing import NDArray
 
     from xarray.core.types import T_Chunks
@@ -631,6 +632,23 @@ def __index__(self, /) -> int:
     def __int__(self, /) -> int:
         return self._data.__int__()
 
+    # dlpack
+    def __dlpack__(
+        self,
+        /,
+        *,
+        stream: int | Any | None = None,
+        max_version: tuple[int, int] | None = None,
+        dl_device: tuple[IntEnum, int] | None = None,
+        copy: bool | None = None,
+    ) -> Any:
+        return self._data.__dlpack__(
+            stream=stream, max_version=max_version, dl_device=dl_device, copy=copy
+        )
+
+    def __dlpack_device__(self, /) -> tuple[IntEnum, int]:
+        return self._data.__dlpack_device__()
+
     # Arithmetic Operators
 
     def __neg__(self, /):