cobra-array 0.1.0__tar.gz → 0.1.1__tar.gz

{cobra_array-0.1.0/src/cobra_array.egg-info → cobra_array-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cobra-array
-Version: 0.1.0
+Version: 0.1.1
 Summary: A backend-agnostic array utility library that unifies array conversion, context control, and cross-library operations across `NumPy`/`PyTorch`-style ecosystems.
 Author-email: Zhen Tian <zhen.tian.cs@gmail.com>
 Project-URL: Homepage, https://github.com/tinchen777/cobra-array.git

{cobra_array-0.1.0 → cobra_array-0.1.1}/src/cobra_array/__init__.py

@@ -91,7 +91,7 @@ from ._utils import (
 )
 
 __author__ = "Zhen Tian"
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 __all__ = [
     "array_spec",

{cobra_array-0.1.0 → cobra_array-0.1.1}/src/cobra_array/_core.py

@@ -22,7 +22,7 @@ from .exceptions import (
 if TYPE_CHECKING:
     from numpy.typing import NDArray
     from .compat import CompatArray
-    from .types import (T, dtypeT, DType, Device, ArrayLike, ArrayLibraryName)
+    from .types import (T, dtypeT, DType, AnyDevice, ArrayLike, ArrayLibraryName)
 
 
 def array_spec(
@@ -85,10 +85,6 @@ def array_spec(
             If `ref` is not `None`, a string, or an integer.
         NotArrayAPIObjectError
             If the reference array determined by `ref` is not an array API compatible array object.
-
-    Examples
-    --------
-    
     """
     kw_arrays = kw_arrays or {}
 
@@ -197,11 +193,11 @@ def context_spec() -> ArraySpec:
 @overload
 def as_context(obj: NDArray[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT, Literal["cpu"]]: ...
 @overload
-def as_context(obj: ArrayLike[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT, Any]: ...
+def as_context(obj: ArrayLike[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT, AnyDevice]: ...
 @overload
-def as_context(obj: ArrayLike[Any], /, *, unify_dtype: Literal[True] = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[Any, Any]: ...
+def as_context(obj: ArrayLike[Any], /, *, unify_dtype: Literal[True] = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[Any, AnyDevice]: ...
 @overload
-def as_context(obj: object, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[False] = ...) -> CompatArray[Any, Any]: ...
+def as_context(obj: object, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[False] = ...) -> CompatArray[Any, AnyDevice]: ...
 @overload
 def as_context(obj: T, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[True]) -> T: ...
 
@@ -236,7 +232,7 @@ def as_context(
 
     Returns
     -------
-        CompatArray[Any, Any]
+        CompatArray[Any, AnyDevice]
             The converted array representation of the object in the current context `compatibility namespace`, with the current context `dtype` and `device` if specified.
         object
             If :param:`arraylike_only` is `True` and the object is not array-like.
@@ -275,7 +271,7 @@ class array_context:
         self,
         xp: Optional[Union[object, ArrayLibraryName]] = None,
         dtype: Optional[DType] = None,
-        device: Optional[Device] = None
+        device: Optional[AnyDevice] = None
     ):
         """
         Initialize the context manager with the specified `array namespace`, `dtype` and `device`.
@@ -286,11 +282,11 @@ class array_context:
                 The target `array namespace` or array library name for the context.
                 - `None`: Use the `compatibility namespace` from the context;
 
-            dtype : Optional[DTypeT], default to `None`
+            dtype : Optional[DType], default to `None`
                 The target `dtype` for the context.
                 - `None`: Use the `dtype` from the context.
 
-            device : Optional[DeviceT], default to `None`
+            device : Optional[AnyDevice], default to `None`
                 The target `device` for the context.
                 - `None`: Use the `device` from the context.
         """

{cobra_array-0.1.0 → cobra_array-0.1.1}/src/cobra_array/compat/_array.py

@@ -103,7 +103,7 @@ class CompatArray(Compat):
 
         Parameters
         ----------
-            device : Device
+            device : AnyDevice
                 A device object or name.
 
             stream : Optional[Union[int, Any]], default to `None`

{cobra_array-0.1.0 → cobra_array-0.1.1}/src/cobra_array/compat/_array.pyi

@@ -9,14 +9,14 @@ from typing import (Union, List, Tuple, Optional, Any, Sequence, Generic, TypeVa
 
 from ._base import Compat
 from ..types import (
-    T, DTypeT, DeviceT, dtypeT, deviceT, DType, Device,
+    T, DTypeT, DeviceT, dtypeT, deviceT, DType, AnyDevice,
     ArrayLike, ArrayLibraryName,
     ArrayOrAny, ArrayOrScalar, ArrayOrReal, ArrayOrIntLike, ArrayOrInt, ArrayOrbool,
     UniqueAllResult, UniqueCountsResult, UniqueInverseResult
 )
 
 TT = TypeVar("TT", bound=DType)
-DT = TypeVar("DT", bound=Device)
+DT = TypeVar("DT", bound=AnyDevice)
 
 
 class CompatArray(Compat, Generic[TT, DT]):
@@ -31,47 +31,54 @@ class CompatArray(Compat, Generic[TT, DT]):
     def from_other(cls, obj: object, /, xp: Literal["numpy"]) -> CompatArray[Any, Literal["cpu"]]: ...
     @overload
     @classmethod
-    def from_other(cls, obj: NDArray[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, Any]: ...
+    def from_other(cls, obj: NDArray[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, AnyDevice]: ...
     @overload
     @classmethod
-    def from_other(cls, obj: ArrayLike[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, Any]: ...
+    def from_other(cls, obj: ArrayLike[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, AnyDevice]: ...
     @overload
     @classmethod
-    def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, Any]: ...
+    def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, AnyDevice]: ...
     @classmethod
-    def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, Any]: ...
+    def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, AnyDevice]: ...
 
     @overload
     def __new__(cls, arr: NDArray[dtypeT], /, **kwargs) -> CompatArray[dtypeT, Literal["cpu"]]: ...
     @overload
-    def __new__(cls, arr: Tensor, /, **kwargs) -> CompatArray[Any, Any]: ...
+    def __new__(cls, arr: Tensor, /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
     @overload
     def __new__(cls, arr: CompatArray[dtypeT, deviceT], /, **kwargs) -> CompatArray[dtypeT, deviceT]: ...
     @overload
-    def __new__(cls, arr: ArrayLike[dtypeT], /, **kwargs) -> CompatArray[dtypeT, Any]: ...
-    def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any, Any]: ...
+    def __new__(cls, arr: ArrayLike[dtypeT], /, **kwargs) -> CompatArray[dtypeT, AnyDevice]: ...
+    def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
 
     # === Conversion functions ===
     def to_numpy(self, *, copy: bool = False) -> NDArray[TT]: ...
-    def to_tensor(self, *, device: Optional[DeviceT] = None, copy: bool = False) -> Tensor: ...
+    def to_tensor(self, *, device: Optional[AnyDevice] = None, copy: bool = False) -> Tensor: ...
     def to_list(self, *, copy: bool = False) -> List[TT]: ...
 
     # === Device functions ===
+    @overload
     def to_device(self, device: DeviceT, /, *, stream: Optional[Any] = None) -> CompatArray[TT, DeviceT]: ...
+    @overload
+    def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[TT, AnyDevice]: ...
+
+    def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[Any, AnyDevice]: ...
 
     # === Data type functions ===
     @overload
-    def astype(self, dtype: DTypeT, /, *, copy: bool = ...) -> CompatArray[DTypeT, DT]: ...
+    def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: None = ...) -> CompatArray[DTypeT, DT]: ...
     @overload
     def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: DeviceT) -> CompatArray[DTypeT, DeviceT]: ...
+    @overload
+    def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: AnyDevice) -> CompatArray[DTypeT, AnyDevice]: ...
 
     def astype(
         self,
         dtype: DType,
         /, *,
         copy: bool = True,
-        device: Optional[Device] = None
-    ) -> CompatArray[Any, Any]:
+        device: Optional[AnyDevice] = None
+    ) -> CompatArray[Any, AnyDevice]:
         """
         Copies `self` to a specified data type irrespective of Type Promotion Rules rules.
 
@@ -88,8 +95,8 @@ class CompatArray(Compat, Generic[TT, DT]):
                     - `True`: `self` must be returned;
                     - `False`: A newly allocated array must be returned.
 
-            device : Optional[DeviceT], default to `None`
-                DeviceT on which to place the returned array.
+            device : Optional[AnyDevice], default to `None`
+                Device on which to place the returned array.
                 - `None`: The output array device must be inferred from `self`.
 
         Returns
@@ -1793,7 +1800,7 @@ class CompatArray(Compat, Generic[TT, DT]):
 
 
 @overload
-def unwrap(obj: CompatArray[dtypeT, Any]) -> ArrayLike[dtypeT]: ...
+def unwrap(obj: CompatArray[dtypeT, AnyDevice]) -> ArrayLike[dtypeT]: ...
 @overload
 def unwrap(obj: ArrayLike[dtypeT]) -> ArrayLike[dtypeT]: ...
 @overload
@@ -1808,7 +1815,7 @@ def wrap_arraylike(arr: NDArray[dtypeT], xp: Optional[object] = ...) -> CompatAr
 @overload
 def wrap_arraylike(arr: CompatArray[dtypeT, deviceT], xp: Optional[object] = ...) -> CompatArray[dtypeT, deviceT]: ...
 @overload
-def wrap_arraylike(arr: ArrayLike[dtypeT], xp: Optional[object] = ...) -> CompatArray[dtypeT, Any]: ...
+def wrap_arraylike(arr: ArrayLike[dtypeT], xp: Optional[object] = ...) -> CompatArray[dtypeT, AnyDevice]: ...
 @overload
 def wrap_arraylike(arr: T, xp: Optional[object] = ...) -> T: ...