cobra-array 0.1.4__tar.gz → 0.2.0__tar.gz

{cobra_array-0.1.4/src/cobra_array.egg-info → cobra_array-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cobra-array
-Version: 0.1.4
+Version: 0.2.0
 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.4 → cobra_array-0.2.0}/src/cobra_array/__init__.py

@@ -32,11 +32,8 @@ Examples
     from cobra_array.convert import to_numpy, to_tensor, to_list
 
     data = [[1, 2], [3, 4]]
-
     arr_np = to_numpy(data, dtype=np.float32)  # numpy.ndarray float32
-
     arr_torch = to_tensor(data, device="cpu")
-
     back_to_list = to_list(arr_np)  # [[1.0, 2.0], [3.0, 4.0]]
 
 - Context-based conversion::
@@ -66,7 +63,6 @@ Examples
     from cobra_array.default import as_default, default_spec
 
     spec = default_spec()
-
     x = as_default([1, 2, 3], unify_dtype=True, unify_device=True)
 """
 
@@ -84,7 +80,7 @@ from ._utils import (
 )
 
 __author__ = "Zhen Tian"
-__version__ = "0.1.4"
+__version__ = "0.2.0"
 
 __all__ = [
     "array_spec",

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/_core.py

@@ -500,9 +500,9 @@ def unify_args(
                 return func(*args, **kwargs)
 
             with array_context.from_array_spec(spec):
-                out_args = tuple(
+                out_args = [
                     as_context(a, unify_dtype=unify_dtype, unify_device=unify_device, arraylike_only=arraylike_only) for a in args
-                )
+                ]
                 out_kwargs = {
                     k: as_context(v, unify_dtype=unify_dtype, unify_device=unify_device, arraylike_only=arraylike_only)
                     for k, v in kwargs.items()

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/_utils.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import array_api_compat as api
 import warnings
 from types import ModuleType
-from typing import Any
+from typing import (Any, overload, Optional, Literal)
 
 from .exceptions import UnsupportedNamespaceError
 
@@ -25,7 +25,13 @@ def warn(msg: str, /, category: Any, stack: int = 2):
     return warnings.warn(msg, category=category, stacklevel=stack+1)
 
 
-def array_namespace_alias(xp: object) -> str:
+@overload
+def array_namespace_alias(xp: object, /, *, raise_on_unsupported: Literal[True] = ...) -> str: ...
+@overload
+def array_namespace_alias(xp: object, /, *, raise_on_unsupported: Literal[False]) -> Optional[str]: ...
+
+
+def array_namespace_alias(xp: object, /, *, raise_on_unsupported: bool = True) -> Optional[str]:
     """
     Get the alias of the `array namespace`.
 
@@ -34,18 +40,23 @@ def array_namespace_alias(xp: object) -> str:
         xp : object
             The `array namespace`.
 
+        raise_on_unsupported : bool, optional
+            Whether to raise an error for unsupported array namespaces.
+
     Returns
     -------
         str
             The alias of the `array namespace`.
             - Including: `"NumPy"`, `"Cupy"`, `"PyTorch"`, `"NDONNX"`, `"Dask"`, `"JAX"`, `"sparse"` and `"array-api-strict"`.
+        `None`
+            If the input is not a supported `array namespace` and `raise_on_unsupported` is `False`.
 
     Raises
     ------
         UnsupportedNameSpaceError
-            If the input object is not a supported `array namespace`.
+            If the input object is not a supported `array namespace` and :param:`raise_on_unsupported` is `True`.
     """
-    if isinstance(xp, ModuleType):
+    if type(xp) is ModuleType:
         if api.is_numpy_namespace(xp):
             return "NumPy"
 
@@ -70,20 +81,17 @@ def array_namespace_alias(xp: object) -> str:
         if api.is_array_api_strict_namespace(xp):
             return "array-api-strict"
 
-    raise UnsupportedNamespaceError(
-        f"Got unsupported array namespace of type {type(xp)}."
-    )
+    if raise_on_unsupported:
+        raise UnsupportedNamespaceError(
+            f"Got unsupported array namespace of type {type(xp)}."
+        )
 
 
 def is_array_namespace(obj: object) -> bool:
     """
     Returns `True` if input is a supported `array namespace`.
     """
-    try:
-        array_namespace_alias(obj)
-        return True
-    except UnsupportedNamespaceError:
-        return False
+    return array_namespace_alias(obj, raise_on_unsupported=False) is not None
 
 
 def is_compat_namespace(xp: object) -> bool:
@@ -91,3 +99,10 @@ def is_compat_namespace(xp: object) -> bool:
     Returns `True` if input is a `compatibility namespace` wrapped by :class:`CompatNamespace`
     """
     return "(compat)" in getattr(xp, "__name__", "")
+
+
+# def is_array_api_object(obj: object) -> bool:
+#     """
+#     Returns `True` if input is an `array API object` that belongs to a supported `array namespace`.
+#     """
+#     return api.is_array_api_obj(obj)

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/compat/__init__.py

@@ -1,6 +1,6 @@
 # src/cobra_array/compat/__init__.py
 """
-Compatibility utilities for :pkg:`cobra_color`.
+Compatibility utilities for :pkg:`cobra_array`.
 
 Functions
 ---------

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/compat/_array.py

@@ -108,16 +108,17 @@ class CompatArray(Compat):
         _xp = to_xp(xp)
         return cls(
             as_array(unwrap(obj), _xp, copy=copy),
-            xp=_xp
+            xp=_xp, check=False
         )
 
     def __new__(cls, arr, /, *, copy=False, **kwargs):
-        if isinstance(arr, CompatArray):
+        if type(arr) is cls:
             # for `CompatArray` input
             return arr.copy() if copy else arr
 
+        _check = kwargs.get("check", True)
         # for non-`CompatArray` input
-        if not api.is_array_api_obj(arr):
+        if _check and not api.is_array_api_obj(arr):
             raise NotArrayAPIObjectError(
                 f"Parameter `arr` of `CompatArray` must be an array API compatible array object, got {type(arr)}."
             )
@@ -135,6 +136,8 @@ class CompatArray(Compat):
         Convert `self` to a `NumPy array`.
         See also :func:`convert.to_numpy`.
         """
+        if self._xp_name == "NumPy" and not copy:
+            return self._arr
         return to_numpy(self._arr, copy=copy)
 
     def to_tensor(self, *, device=None, copy=False):
@@ -142,6 +145,8 @@ class CompatArray(Compat):
         Convert `self` to a `PyTorch tensor`.
         See also :func:`convert.to_tensor`.
         """
+        if self._xp_name == "PyTorch" and not copy and (device is None or device == self.device):
+            return self._arr
         return to_tensor(self._arr, device=device, copy=copy)
 
     def to_list(self, *, copy=False):
@@ -197,7 +202,7 @@ class CompatArray(Compat):
                 All the arrays have the same shape.
         """
         result = self._get_xp_attr("unstack")(self._arr, axis=axis)
-        return tuple(CompatArray(arr, xp=self._xp) for arr in result)
+        return tuple([CompatArray(arr, xp=self._cxp, check=False) for arr in result])
 
     # === Searching functions ===
     def nonzero(self):
@@ -218,7 +223,7 @@ class CompatArray(Compat):
         - If `self` has a boolean data type, non-zero elements are those elements which are equal to `True`.
         """
         result = self._get_xp_attr("nonzero")(self._arr)
-        return tuple(CompatArray(arr, xp=self._xp) for arr in result)
+        return tuple([CompatArray(arr, xp=self._cxp) for arr in result])
 
     # === Set functions ===
     def unique_all(self):
@@ -238,10 +243,10 @@ class CompatArray(Compat):
         """
         result = self._get_xp_attr("unique_all")(self._arr)
         return UniqueResult(
-            values=CompatArray(result.values, xp=self._xp),
-            indices=CompatArray(result.indices, xp=self._xp),
-            inverse_indices=CompatArray(result.inverse_indices, xp=self._xp),
-            counts=CompatArray(result.counts, xp=self._xp),
+            values=CompatArray(result.values, xp=self._cxp, check=False),
+            indices=CompatArray(result.indices, xp=self._cxp, check=False),
+            inverse_indices=CompatArray(result.inverse_indices, xp=self._cxp, check=False),
+            counts=CompatArray(result.counts, xp=self._cxp, check=False),
         )
 
     def unique_counts(self):
@@ -259,10 +264,10 @@ class CompatArray(Compat):
         """
         result = self._get_xp_attr("unique_counts")(self._arr)
         return UniqueResult(
-            values=CompatArray(result.values, xp=self._xp),
+            values=CompatArray(result.values, xp=self._cxp, check=False),
             indices=None,
             inverse_indices=None,
-            counts=CompatArray(result.counts, xp=self._xp),
+            counts=CompatArray(result.counts, xp=self._cxp, check=False),
         )
 
     def unique_inverse(self):
@@ -280,9 +285,9 @@ class CompatArray(Compat):
         """
         result = self._get_xp_attr("unique_inverse")(self._arr)
         return UniqueResult(
-            values=CompatArray(result.values, xp=self._xp),
+            values=CompatArray(result.values, xp=self._cxp, check=False),
             indices=None,
-            inverse_indices=CompatArray(result.inverse_indices, xp=self._xp),
+            inverse_indices=CompatArray(result.inverse_indices, xp=self._cxp, check=False),
             counts=None,
         )
 
@@ -291,7 +296,7 @@ class CompatArray(Compat):
         """
         Return a copy of `self` via :func:`convert.as_array`.
         """
-        return CompatArray.from_other(self._arr, xp=self._xp, copy=True)
+        return CompatArray.from_other(self._arr, xp=self._cxp, copy=True)
 
     def _get_attr(self, name: str):
         """Try to get the attribute `name` from `self`."""
@@ -379,7 +384,7 @@ class CompatArray(Compat):
             result = self._get_xp_attr("T")(self._arr)
         except (AttributeError, TypeError):
             result = self._get_attr("T")
-        return CompatArray(result, xp=self._xp)
+        return CompatArray(result, xp=self._cxp, check=False)
 
     @property
     def mT(self):
@@ -391,19 +396,19 @@ class CompatArray(Compat):
             result = self._get_xp_attr("mT")(self._arr)
         except (AttributeError, TypeError):
             result = self._get_attr("mT")
-        return CompatArray(result, xp=self._xp)
+        return CompatArray(result, xp=self._cxp, check=False)
 
     def __array__(self):
         """Allow implicit NumPy conversion."""
         return self.to_numpy()
 
     def __getattr__(self, name: str):
-        attr = self._get_cxp_attr(name)
+        attr = self._get_xp_attr(name)
 
         if callable(attr) and not isinstance(attr, type):
-            def wrapper(*args, **kwargs):
-                return attr(self._arr, *args, **kwargs)
-            return wrapper
+            wrapped = _make_wrapper(self._xp_name, attr, self._cxp, self._arr)
+            self.__dict__[name] = wrapped
+            return wrapped
         raise CompatArrayAttributeError(f"`CompatArray` `{self._xp_name}` does not support attribute `{name}`.")
 
     def __len__(self):
@@ -540,7 +545,7 @@ def unwrap(obj):
     """
     Unwraps a :class:`CompatArray` array to get the backend-specific array instance, or returns the object itself if it is not a :class:`CompatArray` array.
     """
-    return obj.arr if isinstance(obj, CompatArray) else obj
+    return obj.arr if type(obj) is CompatArray else obj
 
 
 def wrap_arraylike(arr, xp=None):
@@ -549,8 +554,8 @@ def wrap_arraylike(arr, xp=None):
     """
     if api.is_array_api_obj(arr):
         if xp is None:
-            return CompatArray(arr)
-        return CompatArray(arr, xp=xp)
+            return CompatArray(arr, check=False)
+        return CompatArray(arr, xp=xp, check=False)
     return arr
 
 
@@ -559,3 +564,26 @@ def to_cxp(xp):
     from ._namespace import CompatNamespace
 
     return CompatNamespace(xp)
+
+
+def _make_wrapper(xp_name, attr, cxp, first, /):
+    """Make a wrapper function for the attribute `name` of the `array namespace`."""
+    wrap_ = lambda x: wrap_arraylike(x, xp=cxp)
+    unwrap_ = unwrap
+
+    def wrapper(*args, **kwargs):
+        if xp_name == "NumPy":
+            return wrap_(attr(first, *args, **kwargs))
+
+        n = len(args)
+        if n == 0 and not kwargs:
+            return wrap_(attr(first))
+        if n == 1 and not kwargs:
+            return wrap_(attr(first, unwrap_(args[0])))
+        new_args = [unwrap_(a) for a in args]
+        if kwargs:
+            new_kwargs = {k: unwrap_(v) for k, v in kwargs.items()}
+            return wrap_(attr(first, *new_args, **new_kwargs))
+        return wrap_(attr(first, *new_args))
+    wrapper.__name__ = getattr(attr, "__name__", "wrapper")
+    return wrapper

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/compat/_array.pyi

@@ -50,6 +50,8 @@ class CompatArray(Compat, Generic[TT, DT]):
     def __new__(cls, arr: CompatArray[dtypeT, deviceT], /, **kwargs) -> CompatArray[dtypeT, deviceT]: ...
     @overload
     def __new__(cls, arr: ArrayLike[dtypeT], /, **kwargs) -> CompatArray[dtypeT, AnyDevice]: ...
+    @overload
+    def __new__(cls, arr: object, /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
     def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
 
     # === Conversion functions ===

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/compat/_base.py

@@ -3,7 +3,7 @@
 # @TianZhen
 
 from __future__ import annotations
-from typing import (Any, TYPE_CHECKING)
+from typing import (Any, TYPE_CHECKING, Callable)
 
 from .._utils import array_namespace_alias
 

{cobra_array-0.1.4 → cobra_array-0.2.0}/src/cobra_array/compat/_namespace.py

@@ -50,7 +50,7 @@ class CompatNamespace(Compat):
     AttributeError: ...
     """
     def __new__(cls, xp, /):
-        if isinstance(xp, CompatNamespace):
+        if type(xp) is cls:
             # for `CompatNamespace` input
             return xp
         # for `Namespace` input
@@ -91,10 +91,10 @@ class CompatNamespace(Compat):
                 Each returned array should have the same data type as the input arrays.
         """
         result = self._get_xp_attr("meshgrid")(
-            *tuple(unwrap(arr) for arr in arrays),
+            *[unwrap(arr) for arr in arrays],
             indexing=indexing
         )
-        return [CompatArray(arr, xp=self._xp) for arr in result]
+        return [CompatArray(arr, xp=self) for arr in result]
 
     # === Data Type functions ===
     def can_cast(self, from_, to, /):
@@ -170,7 +170,7 @@ class CompatNamespace(Compat):
         """
         return self._get_xp_attr("iinfo")(unwrap(type_))
 
-    def isdtype(self, dtype, kind) -> bool:
+    def isdtype(self, dtype, kind):
         """
         Returns a boolean indicating whether a provided :param:`dtype` is of a specified data type :param:`kind`.
 
@@ -214,7 +214,7 @@ class CompatNamespace(Compat):
             DType
                 The dtype resulting from an operation involving the input arrays, scalars, and/or dtypes.
         """
-        return self._get_xp_attr("result_type")(*tuple(unwrap(arr) for arr in arrays_and_dtypes))
+        return self._get_xp_attr("result_type")(*[unwrap(arr) for arr in arrays_and_dtypes])
 
     # === Manipulation functions ===
     def broadcast_arrays(self, *arrays):
@@ -233,8 +233,8 @@ class CompatNamespace(Compat):
                 Each array must have the same shape.
                 Each array must have the same dtype as its corresponding input array.
         """
-        result = self._get_xp_attr("broadcast_arrays")(*tuple(unwrap(arr) for arr in arrays))
-        return [CompatArray(arr, xp=self._xp) for arr in result]
+        result = self._get_xp_attr("broadcast_arrays")(*[unwrap(arr) for arr in arrays])
+        return [CompatArray(arr, xp=self) for arr in result]
 
     # === Constants ===
     @property
@@ -315,16 +315,37 @@ class CompatNamespace(Compat):
     def __name__(self):
         return "(compat)" + getattr(self._xp, "__name__", type(self._xp).__name__)
 
-    def __getattr__(self, name: str):
+    def __getattr__(self, name):
         attr = self._get_xp_attr(name)
 
         if callable(attr):
-            def wrapper(*args, **kwargs):
-                if not args and not kwargs:
-                    return wrap_arraylike(attr(), xp=self._xp)
-
-                new_args = tuple(unwrap(a) for a in args)
-                new_kwargs = {k: unwrap(v) for k, v in kwargs.items()} if kwargs else kwargs
-                return wrap_arraylike(attr(*new_args, **new_kwargs), xp=self._xp)
-            return wrapper
+            wrapped = _make_wrapper(self._xp_name, attr, self)
+            self.__dict__[name] = wrapped
+            return wrapped
         raise CompatNamespaceAttributeError(f"`CompatNamespace` `{self._xp_name}` does not support attribute `{name}`.")
+
+
+def _make_wrapper(xp_name, attr, cxp, /):
+    """Make a wrapper function for the attribute `name` of the `array namespace`."""
+    wrap_ = lambda x: wrap_arraylike(x, xp=cxp)
+    unwrap_ = unwrap
+
+    def wrapper(*args, **kwargs):
+        if xp_name == "NumPy":
+            return wrap_(attr(*args, **kwargs))
+
+        n = len(args)
+        if n == 0 and not kwargs:
+            return wrap_(attr())
+        if n == 1 and not kwargs:
+            return wrap_(attr(unwrap_(args[0])))
+        if n == 2 and not kwargs:
+            a0, a1 = args
+            return wrap_(attr(unwrap_(a0), unwrap_(a1)))
+        new_args = [unwrap_(a) for a in args]
+        if kwargs:
+            new_kwargs = {k: unwrap_(v) for k, v in kwargs.items()}
+            return wrap_(attr(*new_args, **new_kwargs))
+        return wrap_(attr(*new_args))
+    wrapper.__name__ = getattr(attr, "__name__", "wrapper")
+    return wrapper

{cobra_array-0.1.4 → cobra_array-0.2.0/src/cobra_array.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cobra-array
-Version: 0.1.4
+Version: 0.2.0
 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.2.0/tests/test_compat.py

@@ -0,0 +1,482 @@
+import sys
+from typing import Any
+
+import numpy as np
+import pytest
+import time
+from types import ModuleType
+
+sys.path.insert(0, "/data/tianzhen/my_packages/cobra-array/src")
+
+from cobra_array.array_api import torch_xp
+from cobra_array.compat import CompatArray, unwrap, wrap_arraylike
+from cobra_array.exceptions import CompatArrayAttributeError, NotArrayAPIObjectError
+from cobra_array.compat._array import UniqueResult
+
+
+def _arr_1d():
+    return CompatArray(np.array([1.0, 2.0, 3.0], dtype=np.float32))
+
+
+def _arr_2d():
+    return CompatArray(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))
+
+
+def test_from_other_with_numpy_namespace():
+    out = CompatArray.from_other([1, 2, 3], xp="numpy")
+    assert isinstance(out, CompatArray)
+    assert out.xp_name == "NumPy"
+    assert out.to_list() == [1, 2, 3]
+
+
+def test_new_from_compatarray_copy_and_no_copy():
+    a = _arr_1d()
+    b = CompatArray(a)
+    c = CompatArray(a, copy=True)
+
+    assert b is a
+    assert c is not a
+    assert c.to_list() == a.to_list()
+
+
+def test_new_invalid_input_raises():
+    with pytest.raises(NotArrayAPIObjectError):
+        CompatArray("not-array")  # type: ignore[arg-type]
+
+
+def test_to_numpy_to_list_and_array_protocol():
+    a = _arr_1d()
+
+    n1 = a.to_numpy(copy=False)
+    n2 = np.asarray(a)
+
+    assert isinstance(n1, np.ndarray)
+    assert isinstance(n2, np.ndarray)
+    assert n1.tolist() == [1.0, 2.0, 3.0]
+    assert n2.tolist() == [1.0, 2.0, 3.0]
+
+
+def test_to_device_on_numpy_backend():
+    a = _arr_1d()
+    moved = a.to_device("cpu")
+
+    assert isinstance(moved, np.ndarray)
+    assert moved.tolist() == a.to_list()
+
+
+@pytest.mark.skipif(torch_xp is None, reason="PyTorch not available")
+def test_to_tensor_available_backend():
+    a = _arr_1d()
+    t = a.to_tensor(device="cpu")
+
+    assert isinstance(t, torch_xp.Tensor)
+    assert tuple(t.shape) == (3,)
+
+
+def test_unstack_and_nonzero():
+    a = _arr_2d()
+
+    pieces = a.unstack(axis=0)
+    nz = a.nonzero()
+
+    assert isinstance(pieces, tuple)
+    assert len(pieces) == 2
+    assert all(isinstance(x, CompatArray) for x in pieces)
+    assert pieces[0].to_list() == [1.0, 2.0]
+
+    assert isinstance(nz, tuple)
+    assert len(nz) == 2
+    assert all(isinstance(x, CompatArray) for x in nz)
+
+
+def test_unique_all_unique_counts_unique_inverse():
+    a = CompatArray(np.array([1, 2, 2, 1, 3], dtype=np.int32))
+    
+    print(type(a) is CompatArray)
+    
+    start = time.process_time()
+    all_result = a.cxp.unique_all(a)
+    # counts_result = a.unique_counts()
+    # inverse_result = a.unique_inverse()
+    
+    # print(all_result.values)
+    
+    # assert isinstance(all_result.values, np.ndarray)
+    # assert isinstance(all_result.indices, np.ndarray)
+    # assert isinstance(all_result.inverse_indices, np.ndarray)
+    # assert isinstance(all_result.counts, np.ndarray)
+
+    print(f"(cxp1)Unique operations test took {time.process_time() - start:.5f} seconds")
+    
+
+    start = time.process_time()
+    all_result = a.unique_all()
+    # counts_result = a.unique_counts()
+    # inverse_result = a.unique_inverse()
+    
+    # assert isinstance(all_result.values, CompatArray)
+    # assert isinstance(all_result.indices, CompatArray)
+    # assert isinstance(all_result.inverse_indices, CompatArray)
+    # assert isinstance(all_result.counts, CompatArray)
+
+    print(f"Unique operations test took {time.process_time() - start:.5f} seconds")
+
+    start = time.process_time()
+    all_result = a.cxp.unique_all(a)
+    # counts_result = a.unique_counts()
+    # inverse_result = a.unique_inverse()
+    
+    # print(all_result.values)
+    
+    # assert isinstance(all_result.values, np.ndarray)
+    # assert isinstance(all_result.indices, np.ndarray)
+    # assert isinstance(all_result.inverse_indices, np.ndarray)
+    # assert isinstance(all_result.counts, np.ndarray)
+
+    print(f"(cxp)Unique operations test took {time.process_time() - start:.5f} seconds")
+
+    # assert isinstance(counts_result.values, CompatArray)
+    # assert isinstance(counts_result.counts, CompatArray)
+
+    # assert isinstance(inverse_result.values, CompatArray)
+    # assert isinstance(inverse_result.inverse_indices, CompatArray)
+    
+    a = np.array([1, 2, 2, 1, 3], dtype=np.int32)
+
+    start = time.process_time()
+    all_result = np.unique(a, return_index=True, return_inverse=True, return_counts=True)
+    
+    UniqueResult(
+        CompatArray(all_result[0], xp=np),
+        CompatArray(all_result[1], xp=np),
+        CompatArray(all_result[2], xp=np),
+        CompatArray(all_result[3], xp=np)
+    )
+    
+    print(f"(RAW)Unique operations test took {time.process_time() - start:.5f} seconds")
+    
+    # start = time.process_time()
+    # all_result = torch_xp.unique(torch.tensor(a), return_inverse=True, return_counts=True)
+  
+    # print(f"API operations test took {time.process_time() - start:.4f} seconds")
+    
+    from array_api_compat import get_namespace
+    xp = get_namespace(a)
+
+    start = time.process_time()
+    all_result = xp.unique_all(a)
+    
+    UniqueResult(
+        CompatArray(all_result.values, xp=xp),
+        CompatArray(all_result.indices, xp=xp),
+        CompatArray(all_result.inverse_indices, xp=xp),
+        CompatArray(all_result.counts, xp=xp)
+    )
+    assert isinstance(all_result.values, xp.ndarray)
+    assert isinstance(all_result.indices, xp.ndarray)
+    assert isinstance(all_result.inverse_indices, xp.ndarray)
+    assert isinstance(all_result.counts, xp.ndarray)
+    
+    print(f"API operations test took {time.process_time() - start:.5f} seconds")
+    
+    print(type(xp) is ModuleType)
+
+
+def test_add():
+
+    a = CompatArray(np.array([1, 2, 2, 1, 3], dtype=np.int32))
+    print(type(a) is CompatArray)
+
+    start = time.time()
+    all_result = a.add(1)
+    all_result = a.add(1)
+    # counts_result = a.unique_counts()
+    # inverse_result = a.unique_inverse()
+
+    print(f"add operations test took {time.time() - start:.5f} seconds")
+
+    start = time.time()
+    all_result = a.cxp.add(a, 1)
+    all_result = a.cxp.add(a, 1)
+
+    print(f"(cxp)add operations test took {time.time() - start:.5f} seconds")
+    
+    a = np.array([1, 2, 2, 1, 3], dtype=np.int32)
+
+    start = time.time()
+    all_result = np.add(a, 1)
+
+    print(f"(RAW)add operations test took {time.time() - start:.5f} seconds")
+
+    from array_api_compat import get_namespace
+    xp = get_namespace(a)
+
+    start = time.time()
+    all_result = xp.add(a, 1)
+
+    print(f"API operations test took {time.time() - start:.5f} seconds")
+
+
+def test_copy_and_basic_properties():
+    a = _arr_2d()
+    b = a.copy()
+
+    assert isinstance(a.arr, np.ndarray)
+    assert isinstance(b, CompatArray)
+    assert b is not a
+
+    assert str(a.device) == "cpu"
+    assert a.shape == (2, 2)
+    assert a.ndim == 2
+    assert a.size == 4
+
+
+def test_transpose_properties():
+    a = _arr_2d()
+
+    t = a.T
+
+    assert isinstance(t, CompatArray)
+    assert t.shape == (2, 2)
+
+    with pytest.raises(AttributeError):
+        _ = a.mT
+
+
+def test_len_and_repr():
+    a = _arr_1d()
+    z = CompatArray(np.array(5))
+
+    assert len(a) == 3
+    assert "NumPy_Array" in repr(a)
+
+    with pytest.raises(TypeError):
+        len(z)
+
+
+def test_getitem_and_setitem():
+    a = CompatArray(np.array([10, 20, 30], dtype=np.int32))
+
+    assert int(a[1]) == 20
+
+    a[1] = 99
+    assert a.to_list() == [10, 99, 30]
+
+
+def test_scalar_conversions():
+    b = CompatArray(np.array(True))
+    i = CompatArray(np.array(7))
+    f = CompatArray(np.array(1.5, dtype=np.float32))
+    c = CompatArray(np.array(2 + 3j, dtype=np.complex64))
+
+    assert bool(b) is True
+    assert int(i) == 7
+    assert i.__index__() == 7
+    assert float(f) == pytest.approx(1.5)
+    assert complex(c) == complex(2, 3)
+
+
+def test_numpy_operator_overloads():
+    
+
+    a = np.array([1, 2, 3], dtype=np.int32)
+    b = np.array([3, 2, 1], dtype=np.int32)
+
+    start = time.process_time()
+    assert a.__abs__().tolist() == [1, 2, 3]
+    assert (a + b).tolist() == [4, 4, 4]
+    assert (a - b).tolist() == [-2, 0, 2]
+    assert (a * b).tolist() == [3, 4, 3]
+    assert (a / 2).tolist() == [0.5, 1.0, 1.5]
+    assert (a // 2).tolist() == [0, 1, 1]
+    assert (a % 2).tolist() == [1, 0, 1]
+    assert (a & b).tolist() == [1, 2, 1]
+    assert (a | b).tolist() == [3, 2, 3]
+    assert (a ^ b).tolist() == [2, 0, 2]
+    assert (a << 1).tolist() == [2, 4, 6]
+    assert (a >> 1).tolist() == [0, 1, 1]
+
+    assert (-a).tolist() == [-1, -2, -3]
+    assert (+a).tolist() == [1, 2, 3]
+    assert (~a).tolist() == [-2, -3, -4]
+
+    assert (a == b).tolist() == [False, True, False]
+    assert (a != b).tolist() == [True, False, True]
+    assert (a > b).tolist() == [False, False, True]
+    assert (a >= b).tolist() == [False, True, True]
+    assert (a < b).tolist() == [True, False, False]
+    assert (a <= b).tolist() == [True, True, False]
+
+    print(f"NumPy operator overloads test took {time.process_time() - start:.4f} seconds")
+    
+    a = CompatArray(np.array([1, 2, 3], dtype=np.int32))
+    b = CompatArray(np.array([3, 2, 1], dtype=np.int32))
+    
+    start = time.process_time()
+    assert a.__abs__().to_list() == [1, 2, 3]
+    assert (a + b).to_list() == [4, 4, 4]
+    assert (a - b).to_list() == [-2, 0, 2]
+    assert (a * b).to_list() == [3, 4, 3]
+    assert (a / 2).to_list() == [0.5, 1.0, 1.5]
+    assert (a // 2).to_list() == [0, 1, 1]
+    assert (a % 2).to_list() == [1, 0, 1]
+    assert (a & b).to_list() == [1, 2, 1]
+    assert (a | b).to_list() == [3, 2, 3]
+    assert (a ^ b).to_list() == [2, 0, 2]
+    assert (a << 1).to_list() == [2, 4, 6]
+    assert (a >> 1).to_list() == [0, 1, 1]
+
+    assert (-a).to_list() == [-1, -2, -3]
+    assert (+a).to_list() == [1, 2, 3]
+    assert (~a).to_list() == [-2, -3, -4]
+
+    assert (a == b).to_list() == [False, True, False]
+    assert (a != b).to_list() == [True, False, True]
+    assert (a > b).to_list() == [False, False, True]
+    assert (a >= b).to_list() == [False, True, True]
+    assert (a < b).to_list() == [True, False, False]
+    assert (a <= b).to_list() == [True, True, False]
+
+    print(f"Operator overloads test took {time.process_time() - start:.4f} seconds")
+
+
+# def test_torch_operator_overloads():
+#     a = CompatArray(torch.tensor([1, 2, 3], dtype=torch.int32))
+#     b = CompatArray(torch.tensor([3, 2, 1], dtype=torch.int32))
+
+#     start = time.process_time()
+#     assert a.__abs__().to_list() == [1, 2, 3]
+#     assert (a + b).to_list() == [4, 4, 4]
+#     assert (a - b).to_list() == [-2, 0, 2]
+#     assert (a * b).to_list() == [3, 4, 3]
+#     assert (a / 2).to_list() == [0.5, 1.0, 1.5]
+#     assert (a // 2).to_list() == [0, 1, 1]
+#     assert (a % 2).to_list() == [1, 0, 1]
+#     assert (a & b).to_list() == [1, 2, 1]
+#     assert (a | b).to_list() == [3, 2, 3]
+#     assert (a ^ b).to_list() == [2, 0, 2]
+#     assert (a << 1).to_list() == [2, 4, 6]
+#     assert (a >> 1).to_list() == [0, 1, 1]
+
+#     assert (-a).to_list() == [-1, -2, -3]
+#     assert (+a).to_list() == [1, 2, 3]
+#     assert (~a).to_list() == [-2, -3, -4]
+
+#     assert (a == b).to_list() == [False, True, False]
+#     assert (a != b).to_list() == [True, False, True]
+#     assert (a > b).to_list() == [False, False, True]
+#     assert (a >= b).to_list() == [False, True, True]
+#     assert (a < b).to_list() == [True, False, False]
+#     assert (a <= b).to_list() == [True, True, False]
+
+#     print(f"Operator overloads test took {time.process_time() - start:.4f} seconds")
+
+#     a = torch.tensor([1, 2, 3], dtype=torch.int32)
+#     b = torch.tensor([3, 2, 1], dtype=torch.int32)
+
+#     start = time.process_time()
+#     assert a.__abs__().tolist() == [1, 2, 3]
+#     assert (a + b).tolist() == [4, 4, 4]
+#     assert (a - b).tolist() == [-2, 0, 2]
+#     assert (a * b).tolist() == [3, 4, 3]
+#     assert (a / 2).tolist() == [0.5, 1.0, 1.5]
+#     assert (a // 2).tolist() == [0, 1, 1]
+#     assert (a % 2).tolist() == [1, 0, 1]
+#     assert (a & b).tolist() == [1, 2, 1]
+#     assert (a | b).tolist() == [3, 2, 3]
+#     assert (a ^ b).tolist() == [2, 0, 2]
+#     assert (a << 1).tolist() == [2, 4, 6]
+#     assert (a >> 1).tolist() == [0, 1, 1]
+
+#     assert (-a).tolist() == [-1, -2, -3]
+#     assert (+a).tolist() == [1, 2, 3]
+#     assert (~a).tolist() == [-2, -3, -4]
+
+#     assert (a == b).tolist() == [False, True, False]
+#     assert (a != b).tolist() == [True, False, True]
+#     assert (a > b).tolist() == [False, False, True]
+#     assert (a >= b).tolist() == [False, True, True]
+#     assert (a < b).tolist() == [True, False, False]
+#     assert (a <= b).tolist() == [True, True, False]
+
+#     print(f"Torch operator overloads test took {time.process_time() - start:.4f} seconds")
+
+
+def test_pow_operator_matches_current_implementation_behavior():
+    a = _arr_1d()
+    out = a ** 2
+    assert isinstance(out, CompatArray)
+    assert out.to_list() == [1.0, 4.0, 9.0]
+
+
+def test_matmul_operator():
+    x = CompatArray(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))
+    y = CompatArray(np.array([[2.0, 0.0], [1.0, 2.0]], dtype=np.float32))
+
+    out = x @ y
+
+    assert isinstance(out, CompatArray)
+    assert out.to_list() == [[4.0, 4.0], [10.0, 8.0]]
+
+
+def test_getattr_rejects_non_callable_namespace_attrs():
+    a = _arr_1d()
+    with pytest.raises(CompatArrayAttributeError):
+        _ = getattr(a, "pi")
+
+
+def test_module_helpers_unwrap_wrap_to_cxp():
+    raw = np.array([1, 2, 3])
+    wrapped = wrap_arraylike(raw)
+
+    assert isinstance(wrapped, CompatArray)
+    assert unwrap(wrapped) is raw
+    assert unwrap(raw) is raw
+
+    obj: Any = {"k": 1}
+    assert wrap_arraylike(obj) is obj
+
+    out = CompatArray(raw, xp=np)
+    assert out.xp_name == "NumPy"
+
+
+def test_some_pyi_annotated_methods_via_dynamic_dispatch():
+    a = _arr_2d()
+
+    b = a.astype(np.float64)
+    c = a.abs()
+    d = a.add(1)
+    s = a.sum()
+    r = a.reshape((4,))
+    t = a.take(np.array([0], dtype=np.int32), axis=1)
+
+    assert isinstance(b, CompatArray)
+    assert str(b.dtype) == "float64"
+    assert isinstance(c, CompatArray)
+    assert isinstance(d, CompatArray)
+    assert isinstance(s, CompatArray)
+    assert isinstance(r, CompatArray)
+    assert isinstance(t, CompatArray)
+
+    assert r.shape == (4,)
+    assert t.shape == (2, 1)
+
+
+def test_cxp_of_compatarray_matches_array_namespace():
+    a = _arr_1d()
+    cxp = a.cxp
+
+    assert cxp is not None
+    assert cxp.xp_name == a.xp_name
+
+
+if __name__ == "__main__":
+    test_unique_all_unique_counts_unique_inverse()
+    print("=" * 40)
+    
+    test_numpy_operator_overloads()
+    print("=" * 40)
+    # test_torch_operator_overloads()
+    print("=" * 40)
+
+    test_add()

cobra_array-0.1.4/tests/test_compat.py

@@ -1,261 +0,0 @@
-import sys
-from typing import Any
-
-import numpy as np
-import pytest
-
-sys.path.insert(0, "/data/tianzhen/my_packages/cobra-array/src")
-
-from cobra_array.array_api import torch_xp
-from cobra_array.compat import CompatArray, unwrap, wrap_arraylike
-from cobra_array.exceptions import CompatArrayAttributeError, NotArrayAPIObjectError
-
-
-def _arr_1d():
-    return CompatArray(np.array([1.0, 2.0, 3.0], dtype=np.float32))
-
-
-def _arr_2d():
-    return CompatArray(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))
-
-
-def test_from_other_with_numpy_namespace():
-    out = CompatArray.from_other([1, 2, 3], xp="numpy")
-    assert isinstance(out, CompatArray)
-    assert out.xp_name == "NumPy"
-    assert out.to_list() == [1, 2, 3]
-
-
-def test_new_from_compatarray_copy_and_no_copy():
-    a = _arr_1d()
-    b = CompatArray(a)
-    c = CompatArray(a, copy=True)
-
-    assert b is a
-    assert c is not a
-    assert c.to_list() == a.to_list()
-
-
-def test_new_invalid_input_raises():
-    with pytest.raises(NotArrayAPIObjectError):
-        CompatArray("not-array")  # type: ignore[arg-type]
-
-
-def test_to_numpy_to_list_and_array_protocol():
-    a = _arr_1d()
-
-    n1 = a.to_numpy(copy=False)
-    n2 = np.asarray(a)
-
-    assert isinstance(n1, np.ndarray)
-    assert isinstance(n2, np.ndarray)
-    assert n1.tolist() == [1.0, 2.0, 3.0]
-    assert n2.tolist() == [1.0, 2.0, 3.0]
-
-
-def test_to_device_on_numpy_backend():
-    a = _arr_1d()
-    moved = a.to_device("cpu")
-
-    assert isinstance(moved, np.ndarray)
-    assert moved.tolist() == a.to_list()
-
-
-@pytest.mark.skipif(torch_xp is None, reason="PyTorch not available")
-def test_to_tensor_available_backend():
-    a = _arr_1d()
-    t = a.to_tensor(device="cpu")
-
-    assert isinstance(t, torch_xp.Tensor)
-    assert tuple(t.shape) == (3,)
-
-
-def test_unstack_and_nonzero():
-    a = _arr_2d()
-
-    pieces = a.unstack(axis=0)
-    nz = a.nonzero()
-
-    assert isinstance(pieces, tuple)
-    assert len(pieces) == 2
-    assert all(isinstance(x, CompatArray) for x in pieces)
-    assert pieces[0].to_list() == [1.0, 2.0]
-
-    assert isinstance(nz, tuple)
-    assert len(nz) == 2
-    assert all(isinstance(x, CompatArray) for x in nz)
-
-
-def test_unique_all_unique_counts_unique_inverse():
-    a = CompatArray(np.array([1, 2, 2, 1, 3], dtype=np.int32))
-
-    all_result = a.unique_all()
-    counts_result = a.unique_counts()
-    inverse_result = a.unique_inverse()
-
-    assert isinstance(all_result.values, CompatArray)
-    assert isinstance(all_result.indices, CompatArray)
-    assert isinstance(all_result.inverse_indices, CompatArray)
-    assert isinstance(all_result.counts, CompatArray)
-
-    assert isinstance(counts_result.values, CompatArray)
-    assert isinstance(counts_result.counts, CompatArray)
-
-    assert isinstance(inverse_result.values, CompatArray)
-    assert isinstance(inverse_result.inverse_indices, CompatArray)
-
-
-def test_copy_and_basic_properties():
-    a = _arr_2d()
-    b = a.copy()
-
-    assert isinstance(a.arr, np.ndarray)
-    assert isinstance(b, CompatArray)
-    assert b is not a
-
-    assert str(a.device) == "cpu"
-    assert a.shape == (2, 2)
-    assert a.ndim == 2
-    assert a.size == 4
-
-
-def test_transpose_properties():
-    a = _arr_2d()
-
-    t = a.T
-
-    assert isinstance(t, CompatArray)
-    assert t.shape == (2, 2)
-
-    with pytest.raises(AttributeError):
-        _ = a.mT
-
-
-def test_len_and_repr():
-    a = _arr_1d()
-    z = CompatArray(np.array(5))
-
-    assert len(a) == 3
-    assert "NumPy_Array" in repr(a)
-
-    with pytest.raises(TypeError):
-        len(z)
-
-
-def test_getitem_and_setitem():
-    a = CompatArray(np.array([10, 20, 30], dtype=np.int32))
-
-    assert int(a[1]) == 20
-
-    a[1] = 99
-    assert a.to_list() == [10, 99, 30]
-
-
-def test_scalar_conversions():
-    b = CompatArray(np.array(True))
-    i = CompatArray(np.array(7))
-    f = CompatArray(np.array(1.5, dtype=np.float32))
-    c = CompatArray(np.array(2 + 3j, dtype=np.complex64))
-
-    assert bool(b) is True
-    assert int(i) == 7
-    assert i.__index__() == 7
-    assert float(f) == pytest.approx(1.5)
-    assert complex(c) == complex(2, 3)
-
-
-def test_operator_overloads():
-    a = CompatArray(np.array([1, 2, 3], dtype=np.int32))
-    b = CompatArray(np.array([3, 2, 1], dtype=np.int32))
-
-    assert a.__abs__().to_list() == [1, 2, 3]
-    assert (a + b).to_list() == [4, 4, 4]
-    assert (a - b).to_list() == [-2, 0, 2]
-    assert (a * b).to_list() == [3, 4, 3]
-    assert (a / 2).to_list() == [0.5, 1.0, 1.5]
-    assert (a // 2).to_list() == [0, 1, 1]
-    assert (a % 2).to_list() == [1, 0, 1]
-    assert (a & b).to_list() == [1, 2, 1]
-    assert (a | b).to_list() == [3, 2, 3]
-    assert (a ^ b).to_list() == [2, 0, 2]
-    assert (a << 1).to_list() == [2, 4, 6]
-    assert (a >> 1).to_list() == [0, 1, 1]
-
-    assert (-a).to_list() == [-1, -2, -3]
-    assert (+a).to_list() == [1, 2, 3]
-    assert (~a).to_list() == [-2, -3, -4]
-
-    assert (a == b).to_list() == [False, True, False]
-    assert (a != b).to_list() == [True, False, True]
-    assert (a > b).to_list() == [False, False, True]
-    assert (a >= b).to_list() == [False, True, True]
-    assert (a < b).to_list() == [True, False, False]
-    assert (a <= b).to_list() == [True, True, False]
-
-
-def test_pow_operator_matches_current_implementation_behavior():
-    a = _arr_1d()
-    out = a ** 2
-    assert isinstance(out, CompatArray)
-    assert out.to_list() == [1.0, 4.0, 9.0]
-
-
-def test_matmul_operator():
-    x = CompatArray(np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32))
-    y = CompatArray(np.array([[2.0, 0.0], [1.0, 2.0]], dtype=np.float32))
-
-    out = x @ y
-
-    assert isinstance(out, CompatArray)
-    assert out.to_list() == [[4.0, 4.0], [10.0, 8.0]]
-
-
-def test_getattr_rejects_non_callable_namespace_attrs():
-    a = _arr_1d()
-    with pytest.raises(CompatArrayAttributeError):
-        _ = getattr(a, "pi")
-
-
-def test_module_helpers_unwrap_wrap_to_cxp():
-    raw = np.array([1, 2, 3])
-    wrapped = wrap_arraylike(raw)
-
-    assert isinstance(wrapped, CompatArray)
-    assert unwrap(wrapped) is raw
-    assert unwrap(raw) is raw
-
-    obj: Any = {"k": 1}
-    assert wrap_arraylike(obj) is obj
-
-    out = CompatArray(raw, xp=np)
-    assert out.xp_name == "NumPy"
-
-
-def test_some_pyi_annotated_methods_via_dynamic_dispatch():
-    a = _arr_2d()
-
-    b = a.astype(np.float64)
-    c = a.abs()
-    d = a.add(1)
-    s = a.sum()
-    r = a.reshape((4,))
-    t = a.take(np.array([0], dtype=np.int32), axis=1)
-
-    assert isinstance(b, CompatArray)
-    assert str(b.dtype) == "float64"
-    assert isinstance(c, CompatArray)
-    assert isinstance(d, CompatArray)
-    assert isinstance(s, CompatArray)
-    assert isinstance(r, CompatArray)
-    assert isinstance(t, CompatArray)
-
-    assert r.shape == (4,)
-    assert t.shape == (2, 1)
-
-
-def test_cxp_of_compatarray_matches_array_namespace():
-    a = _arr_1d()
-    cxp = a.cxp
-
-    assert cxp is not None
-    assert cxp.xp_name == a.xp_name