Update a function signature from the array API spec

2 files changed, 6 insertions, 5 deletions

diff --git a/numpy/_array_api/_array_object.py b/numpy/_array_api/_array_object.py
index 30858b7c5..5f169ab66 100644
--- a/numpy/_array_api/_array_object.py
+++ b/numpy/_array_api/_array_object.py
@@ -22,7 +22,7 @@ from ._dtypes import _boolean_dtypes, _integer_dtypes, _floating_dtypes
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from ._types import Optional, PyCapsule, Tuple, Union, array
+    from ._types import Any, Optional, PyCapsule, Tuple, Union, array
 
 import numpy as np
 
@@ -186,7 +186,7 @@ class ndarray:
         res = self._array.__bool__()
         return res
 
-    def __dlpack__(self: array, /, *, stream: Optional[int] = None) -> PyCapsule:
+    def __dlpack__(self: array, /, *, stream: Optional[Union[int, Any]] = None) -> PyCapsule:
         """
         Performs the operation __dlpack__.
         """
diff --git a/numpy/_array_api/_types.py b/numpy/_array_api/_types.py
index 36c9aa610..1086699fc 100644
--- a/numpy/_array_api/_types.py
+++ b/numpy/_array_api/_types.py
@@ -6,10 +6,11 @@ annotations in the function signatures. The functions in the module are only
 valid for inputs that match the given type annotations.
 """
 
-__all__ = ['List', 'Literal', 'Optional', 'Tuple', 'Union', 'array', 'device',
-           'dtype', 'SupportsDLPack', 'SupportsBufferProtocol', 'PyCapsule']
+__all__ = ['Any', 'List', 'Literal', 'Optional', 'Tuple', 'Union', 'array',
+           'device', 'dtype', 'SupportsDLPack', 'SupportsBufferProtocol',
+           'PyCapsule']
 
-from typing import List, Literal, Optional, Tuple, Union, TypeVar
+from typing import Any, List, Literal, Optional, Tuple, Union, TypeVar
 
 from . import (ndarray, int8, int16, int32, int64, uint8, uint16, uint32,
                uint64, float32, float64)