cobra-array 0.1.5__tar.gz → 0.2.1__tar.gz

{cobra_array-0.1.5/src/cobra_array.egg-info → cobra_array-0.2.1}/PKG-INFO

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

@@ -80,7 +80,7 @@ from ._utils import (
 )
 
 __author__ = "Zhen Tian"
-__version__ = "0.1.5"
+__version__ = "0.2.1"
 
 __all__ = [
     "array_spec",

{cobra_array-0.1.5 → cobra_array-0.2.1}/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.5 → cobra_array-0.2.1}/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.5 → cobra_array-0.2.1}/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`."""
@@ -334,7 +339,7 @@ class CompatArray(Compat):
     @property
     def device(self):
         """
-        DeviceT on which `self` is stored.
+        Device on which `self` is stored.
         """
         return api.device(self._arr)
 
@@ -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)
+    def __getattr__(self, 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.5 → cobra_array-0.2.1}/src/cobra_array/compat/_array.pyi

@@ -10,7 +10,7 @@ from typing import (Union, List, Tuple, Optional, Any, Sequence, Generic, TypeVa
 from ._base import Compat
 from ._namespace import CompatNamespace
 from ..types import (
-    T, DTypeT, DeviceT, dtypeT, deviceT, DType, AnyDevice,
+    T, dtypeT, DTypeT, deviceT, DeviceT, DType, AnyDevice,
     ArrayLike, ArrayLibraryName,
     ArrayOrAny, ArrayOrScalar, ArrayOrReal, ArrayOrIntLike, ArrayOrInt, ArrayOrbool,
     UniqueAllResult, UniqueCountsResult, UniqueInverseResult
@@ -64,7 +64,6 @@ class CompatArray(Compat, Generic[TT, DT]):
     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 ===
@@ -1771,6 +1770,8 @@ class CompatArray(Compat, Generic[TT, DT]):
 
     def __array__(self) -> NDArray[TT]: ...
 
+    def __getattr__(self, name: str) -> Any: ...
+
     def __len__(self) -> int: ...
     def __abs__(self) -> CompatArray[TT, DT]: ...
     def __add__(self, other: ArrayOrScalar, /) -> CompatArray[Any, DT]: ...

{cobra_array-0.1.5 → cobra_array-0.2.1}/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.5 → cobra_array-0.2.1}/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,139 @@ 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]
+
+    # === Linear Algebra Extension ===
+    def vector_norm(self, x, /, *, axis=None, keepdims=False, ord=2):
+        """
+        Computes the vector norm of a vector (or batch of vectors) :param:`x`.
+
+        Parameters
+        ----------
+            x : ArrayLike[Any]
+                The input array. Should have a floating-point data type.
+
+            axis : Optional[Union[int, Tuple[int, ...]]], default to `None`
+                - _int_: :param:`axis` specifies the axis (dimension) along which to compute vector norms;
+                - _tuple_: :param:`axis` specifies the axes (dimensions) along which to compute batched vector norms;
+                - `None`: The vector norm must be computed over all array values (i.e., equivalent to computing the vector norm of a flattened array).
+
+                Negative indices must be supported.
+
+            keepdims : bool, default to `False`
+                - `True`: The axes (dimensions) specified by :param:`axis` must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see Broadcasting);
+                - `False`: The axes (dimensions) specified by :param:`axis` must not be included in the result.
+
+            ord : Union[int, float, Literal['inf', '-inf']], default to `2`
+                Order of the norm.
+                The following mathematical norms must be supported:
+                - `1`: L1-norm (Manhattan);
+                - `2`: L2-norm (Euclidean);
+                - `"inf"`: infinity norm;
+                - _int_ or _float_ (>=1): p-norm.
+
+                The following non-mathematical “norms” must be supported:
+                - `0`: sum(a != 0);
+                - `-1`: 1./sum(1./abs(a));
+                - `-2`: 1./sqrt(sum(1./a**2));
+                - `"-inf"`: min(abs(a));
+                - _int_ or _float_ (<1): sum(abs(a)**ord)**(1./ord).
+
+        Returns
+        -------
+            CompatArray
+                A :class:`CompatArray` array containing the vector norms.
+                - :param:`axis` is `None`: The returned array must be a zero-dimensional array containing a vector norm;
+                - :param:`axis` is a scalar value (_int_ or _float_): The returned array must have a rank which is one less than the rank of :param:`x`;
+                - :param:`axis` is _tuple_ (`n` elements): The returned array must have a rank which is `n` less than the rank of :param:`x`;
+
+                - :param:`x` is real-valued data type: The returned array must have a real-valued floating-point data type determined by Type Promotion Rules;
+                - :param:`x` is complex-valued data type: The returned array must have a real-valued floating-point data type whose precision matches the precision of :param:`x` (e.g., if :param:`x` is complex128, then the returned array must have a float64 data type).
+        """
+        if ord == "inf":
+            ord = float("inf")
+        elif ord == "-inf":
+            ord = float("-inf")
+
+        result = getattr(self.linalg, "vector_norm")(unwrap(x), axis=axis, keepdims=keepdims, ord=ord)
+        return CompatArray(result, xp=self)
+
+    def matrix_norm(self, x, /, *, keepdims=False, ord="fro"):
+        """
+        Computes the matrix norm of a matrix (or a stack of matrices) :param:`x`.
+
+        Parameters
+        ----------
+            x : ArrayLike[Any]
+                Input array having shape (..., `M`, `N`) and whose innermost two dimensions form `MxN` matrices. Should have a floating-point data type.
+
+            keepdims : bool, default to `False`
+                - `True`: The last two axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see Broadcasting);
+                - `False`: The last two axes (dimensions) must not be included in the result.
+
+            ord : Optional[Union[int, float, Literal['inf', '-inf', 'fro', 'nuc']]], default to `"fro"`
+                order of the norm.
+                The following mathematical norms must be supported:
+                - `"fro"`: Frobenius norm;
+                - `"nuc"`: nuclear norm;
+                - `1`: max(sum(abs(x), axis=0)). The norm corresponds to the induced matrix norm where `p=1` (i.e., the maximum absolute value column sum);
+                - `2`: largest singular value. The norm corresponds to the induced matrix norm where `p=inf` (i.e., the maximum absolute value row sum);
+                - `"inf"`: max(sum(abs(x), axis=1)). The norm corresponds to the induced matrix norm where `p=2` (i.e., the largest singular value).
+
+                The following non-mathematical “norms” must be supported:
+                - `-1`: min(sum(abs(x), axis=0));
+                - `-2`: smallest singular value;
+                - `"-inf"`: min(sum(abs(x), axis=1)).
+
+        Returns
+        -------
+            CompatArray
+                A :class:`CompatArray` array containing the norms for each `MxN` matrix.
+                - :param:`keepdims` is `False`: The returned array must have a rank which is two less than the rank of :param:`x`;
+
+                - :param:`x` is real-valued data type: The returned array must have a real-valued floating-point data type determined by Type Promotion Rules;
+                - :param:`x` is complex-valued data type: The returned array must have a real-valued floating-point data type whose precision matches the precision of :param:`x` (e.g., if :param:`x` is complex128, then the returned array must have a float64 data type).
+
+        """
+        if ord == "inf":
+            ord = float("inf")
+        elif ord == "-inf":
+            ord = float("-inf")
+
+        result = getattr(self.linalg, "matrix_norm")(unwrap(x), keepdims=keepdims, ord=ord)
+        return CompatArray(result, xp=self)
+
+    @property
+    def linalg(self):
+        """
+        The `linalg` namespace for linear algebra functions.
+        The following functions must be supported in the `linalg` namespace:
+        - `cholesky`(x, /, *, upper=False): Returns the lower (upper) Cholesky decomposition of a complex Hermitian or real symmetric positive-definite matrix x.
+        - `cross`(x1, x2, /, *, axis=-1): Returns the cross product of 3-element vectors.
+        - `det`(x, /): Returns the determinant of a square matrix (or a stack of square matrices) x.
+        - `diagonal`(x, /, *, offset=0): Returns the specified diagonals of a matrix (or a stack of matrices) x.
+        - `eigh`(x, /): Returns an eigenvalue decomposition of a complex Hermitian or real symmetric matrix (or a stack of matrices) x.
+        - `eigvalsh`(x, /): Returns the eigenvalues of a complex Hermitian or real symmetric matrix (or a stack of matrices) x.
+        - `inv`(x, /): Returns the multiplicative inverse of a square matrix (or a stack of square matrices) x.
+        - `matmul`(x1, x2, /): Alias for matmul().
+        - `matrix_norm`(x, /, *, keepdims=False, ord='fro'): Computes the matrix norm of a matrix (or a stack of matrices) x.
+        - `matrix_power`(x, n, /): Raises a square matrix (or a stack of square matrices) x to an integer power n.
+        - `matrix_rank`(x, /, *, rtol=None): Returns the rank (i.e., number of non-zero singular values) of a matrix (or a stack of matrices).
+        - `matrix_transpose`(x, /): Alias for matrix_transpose().
+        - `outer`(x1, x2, /): Returns the outer product of two vectors x1 and x2.
+        - `pinv`(x, /, *, rtol=None): Returns the (Moore-Penrose) pseudo-inverse of a matrix (or a stack of matrices) x.
+        - `qr`(x, /, *, mode='reduced'): Returns the QR decomposition of a full column rank matrix (or a stack of matrices).
+        - `slogdet`(x, /): Returns the sign and the natural logarithm of the absolute value of the determinant of a square matrix (or a stack of square matrices) x.
+        - `solve`(x1, x2, /): Returns the solution of a square system of linear equations with a unique solution.
+        - `svd`(x, /, *, full_matrices=True): Returns a singular value decomposition (SVD) of a matrix (or a stack of matrices) x.
+        - `svdvals`(x, /): Returns the singular values of a matrix (or a stack of matrices) x.
+        - `tensordot`(x1, x2, /, *, axes=2): Alias for tensordot().
+        - `trace`(x, /, *, offset=0, dtype=None): Returns the sum along the specified diagonals of a matrix (or a stack of matrices) x.
+        - `vecdot`(x1, x2, /, *, axis=-1): Alias for vecdot().
+        - `vector_norm`(x, /, *, axis=None, keepdims=False, ord=2): Computes the vector norm of a vector (or batch of vectors) x.
+        """
+        return self._get_xp_attr("linalg")
 
     # === Constants ===
     @property
@@ -315,16 +446,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.5 → cobra_array-0.2.1}/src/cobra_array/compat/_namespace.pyi

@@ -4,12 +4,13 @@
 
 from __future__ import annotations
 from numpy.typing import NDArray
+from types import ModuleType
 from typing import (Union, List, Tuple, Optional, Any, Literal, overload)
 
 from ._base import Compat
 from ._array import CompatArray
 from ..types import (
-    DTypeT, DeviceT, dtypeT, DType, AnyDevice,
+    dtypeT, DTypeT, deviceT, DeviceT, DType, AnyDevice,
     ValueT, Value, ArrayLike, ArrayOrAny
 )
 
@@ -802,7 +803,7 @@ class CompatNamespace(Compat):
         Returns
         -------
             CompatArray
-                A :class:`CompatArray` output array containing the concatenated values. 
+                A :class:`CompatArray` output array containing the concatenated values.
         """
         ...
 
@@ -841,6 +842,26 @@ class CompatNamespace(Compat):
         """
         ...
 
+    # === Linear Algebra Extension ===
+    @overload
+    def vector_norm(self, x: NDArray[Any], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = ..., keepdims: bool = ..., ord: Union[int, float, Literal["inf", "-inf"]] = ...) -> CompatArray[float, Literal["cpu"]]: ...
+    @overload
+    def vector_norm(self, x: CompatArray[Any, deviceT], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = ..., keepdims: bool = ..., ord: Union[int, float, Literal["inf", "-inf"]] = ...) -> CompatArray[float, deviceT]: ...
+    @overload
+    def vector_norm(self, x: ArrayLike[Any], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = ..., keepdims: bool = ..., ord: Union[int, float, Literal["inf", "-inf"]] = ...) -> CompatArray[float, AnyDevice]: ...
+    def vector_norm(self, x: ArrayLike[Any], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Union[int, float, Literal["inf", "-inf"]] = 2) -> CompatArray[float, AnyDevice]: ...
+
+    @overload
+    def matrix_norm(self, x: NDArray[Any], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, Literal["cpu"]]: ...
+    @overload
+    def matrix_norm(self, x: CompatArray[Any, deviceT], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, deviceT]: ...
+    @overload
+    def matrix_norm(self, x: ArrayLike[Any], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, AnyDevice]: ...
+    def matrix_norm(self, x: ArrayLike[Any], /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = "fro") -> CompatArray[float, AnyDevice]: ...
+
+    @property
+    def linalg(self) -> ModuleType: ...
+
     # === Constants ===
     @property
     def e(self) -> float: ...
@@ -883,3 +904,5 @@ class CompatNamespace(Compat):
 
     @property
     def __name__(self) -> str: ...
+
+    def __getattr__(self, name: str) -> Any: ...

{cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/convert.py

@@ -122,11 +122,11 @@ def to_tensor(obj, /, *, dtype=None, device=None, copy=True):
             - `None`: Raises `ConvertNoneTypeError`;
             - _others_: Converted to a `PyTorch tensor` directly.
 
-        dtype : Optional[DTypeT], default to `None`
+        dtype : Optional[DType], default to `None`
             The data type of the resulting `PyTorch tensor`.
             - `None`: Use the default data type of the object.
 
-        device : Optional[DeviceT], default to `None`
+        device : Optional[AnyDevice], default to `None`
             The device on which the resulting `PyTorch tensor` will be allocated.
             - `None`: Use the default device (usually `"cpu"`).
 
@@ -294,11 +294,11 @@ def as_array(obj, xp, /, *, dtype=None, device=None, copy=False, arraylike_only=
             - _ArrayLibraryName_ (`"numpy"` or `"torch"`): Converted to a `NumPy array` or `PyTorch tensor` respectively using the corresponding conversion functions;
             - _Namespace_ or _CompatNamespace_: Converted to an array using the `asarray()` function provided by the namespace module, which must be compatible with the array API standard.
 
-        dtype : Optional[DTypeT], default to `None`
+        dtype : Optional[DType], default to `None`
             The data type of the resulting array.
             - `None`: Use the default data type of the object.
 
-        device : Optional[DeviceT], default to `None`
+        device : Optional[AnyDevice], default to `None`
             The device on which the resulting array will be allocated (only if `array namespace` supports it).
 
         copy : bool, default to `False`