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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (32) hide show
  1. {cobra_array-0.1.5/src/cobra_array.egg-info → cobra_array-0.2.1}/PKG-INFO +1 -1
  2. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/__init__.py +1 -1
  3. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/_core.py +2 -2
  4. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/_utils.py +27 -12
  5. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_array.py +53 -25
  6. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_array.pyi +3 -2
  7. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_base.py +1 -1
  8. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_namespace.py +168 -16
  9. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_namespace.pyi +25 -2
  10. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/convert.py +4 -4
  11. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/types.py +3 -2
  12. {cobra_array-0.1.5 → cobra_array-0.2.1/src/cobra_array.egg-info}/PKG-INFO +1 -1
  13. cobra_array-0.2.1/tests/test_compat.py +523 -0
  14. cobra_array-0.1.5/tests/test_compat.py +0 -261
  15. {cobra_array-0.1.5 → cobra_array-0.2.1}/LICENSE +0 -0
  16. {cobra_array-0.1.5 → cobra_array-0.2.1}/README.md +0 -0
  17. {cobra_array-0.1.5 → cobra_array-0.2.1}/pyproject.toml +0 -0
  18. {cobra_array-0.1.5 → cobra_array-0.2.1}/setup.cfg +0 -0
  19. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/array_api.py +0 -0
  20. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/__init__.py +0 -0
  21. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/convert.pyi +0 -0
  22. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/default.py +0 -0
  23. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/exceptions.py +0 -0
  24. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/SOURCES.txt +0 -0
  25. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/dependency_links.txt +0 -0
  26. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/requires.txt +0 -0
  27. {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/top_level.txt +0 -0
  28. {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_backend.py +0 -0
  29. {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_compat_namespace.py +0 -0
  30. {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_convert.py +0 -0
  31. {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_default.py +0 -0
  32. {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_wrap.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cobra-array
3
- Version: 0.1.5
3
+ Version: 0.2.1
4
4
  Summary: A backend-agnostic array utility library that unifies array conversion, context control, and cross-library operations across `NumPy`/`PyTorch`-style ecosystems.
5
5
  Author-email: Zhen Tian <zhen.tian.cs@gmail.com>
6
6
  Project-URL: Homepage, https://github.com/tinchen777/cobra-array.git
@@ -80,7 +80,7 @@ from ._utils import (
80
80
  )
81
81
 
82
82
  __author__ = "Zhen Tian"
83
- __version__ = "0.1.5"
83
+ __version__ = "0.2.1"
84
84
 
85
85
  __all__ = [
86
86
  "array_spec",
@@ -500,9 +500,9 @@ def unify_args(
500
500
  return func(*args, **kwargs)
501
501
 
502
502
  with array_context.from_array_spec(spec):
503
- out_args = tuple(
503
+ out_args = [
504
504
  as_context(a, unify_dtype=unify_dtype, unify_device=unify_device, arraylike_only=arraylike_only) for a in args
505
- )
505
+ ]
506
506
  out_kwargs = {
507
507
  k: as_context(v, unify_dtype=unify_dtype, unify_device=unify_device, arraylike_only=arraylike_only)
508
508
  for k, v in kwargs.items()
@@ -6,7 +6,7 @@ from __future__ import annotations
6
6
  import array_api_compat as api
7
7
  import warnings
8
8
  from types import ModuleType
9
- from typing import Any
9
+ from typing import (Any, overload, Optional, Literal)
10
10
 
11
11
  from .exceptions import UnsupportedNamespaceError
12
12
 
@@ -25,7 +25,13 @@ def warn(msg: str, /, category: Any, stack: int = 2):
25
25
  return warnings.warn(msg, category=category, stacklevel=stack+1)
26
26
 
27
27
 
28
- def array_namespace_alias(xp: object) -> str:
28
+ @overload
29
+ def array_namespace_alias(xp: object, /, *, raise_on_unsupported: Literal[True] = ...) -> str: ...
30
+ @overload
31
+ def array_namespace_alias(xp: object, /, *, raise_on_unsupported: Literal[False]) -> Optional[str]: ...
32
+
33
+
34
+ def array_namespace_alias(xp: object, /, *, raise_on_unsupported: bool = True) -> Optional[str]:
29
35
  """
30
36
  Get the alias of the `array namespace`.
31
37
 
@@ -34,18 +40,23 @@ def array_namespace_alias(xp: object) -> str:
34
40
  xp : object
35
41
  The `array namespace`.
36
42
 
43
+ raise_on_unsupported : bool, optional
44
+ Whether to raise an error for unsupported array namespaces.
45
+
37
46
  Returns
38
47
  -------
39
48
  str
40
49
  The alias of the `array namespace`.
41
50
  - Including: `"NumPy"`, `"Cupy"`, `"PyTorch"`, `"NDONNX"`, `"Dask"`, `"JAX"`, `"sparse"` and `"array-api-strict"`.
51
+ `None`
52
+ If the input is not a supported `array namespace` and `raise_on_unsupported` is `False`.
42
53
 
43
54
  Raises
44
55
  ------
45
56
  UnsupportedNameSpaceError
46
- If the input object is not a supported `array namespace`.
57
+ If the input object is not a supported `array namespace` and :param:`raise_on_unsupported` is `True`.
47
58
  """
48
- if isinstance(xp, ModuleType):
59
+ if type(xp) is ModuleType:
49
60
  if api.is_numpy_namespace(xp):
50
61
  return "NumPy"
51
62
 
@@ -70,20 +81,17 @@ def array_namespace_alias(xp: object) -> str:
70
81
  if api.is_array_api_strict_namespace(xp):
71
82
  return "array-api-strict"
72
83
 
73
- raise UnsupportedNamespaceError(
74
- f"Got unsupported array namespace of type {type(xp)}."
75
- )
84
+ if raise_on_unsupported:
85
+ raise UnsupportedNamespaceError(
86
+ f"Got unsupported array namespace of type {type(xp)}."
87
+ )
76
88
 
77
89
 
78
90
  def is_array_namespace(obj: object) -> bool:
79
91
  """
80
92
  Returns `True` if input is a supported `array namespace`.
81
93
  """
82
- try:
83
- array_namespace_alias(obj)
84
- return True
85
- except UnsupportedNamespaceError:
86
- return False
94
+ return array_namespace_alias(obj, raise_on_unsupported=False) is not None
87
95
 
88
96
 
89
97
  def is_compat_namespace(xp: object) -> bool:
@@ -91,3 +99,10 @@ def is_compat_namespace(xp: object) -> bool:
91
99
  Returns `True` if input is a `compatibility namespace` wrapped by :class:`CompatNamespace`
92
100
  """
93
101
  return "(compat)" in getattr(xp, "__name__", "")
102
+
103
+
104
+ # def is_array_api_object(obj: object) -> bool:
105
+ # """
106
+ # Returns `True` if input is an `array API object` that belongs to a supported `array namespace`.
107
+ # """
108
+ # return api.is_array_api_obj(obj)
@@ -108,16 +108,17 @@ class CompatArray(Compat):
108
108
  _xp = to_xp(xp)
109
109
  return cls(
110
110
  as_array(unwrap(obj), _xp, copy=copy),
111
- xp=_xp
111
+ xp=_xp, check=False
112
112
  )
113
113
 
114
114
  def __new__(cls, arr, /, *, copy=False, **kwargs):
115
- if isinstance(arr, CompatArray):
115
+ if type(arr) is cls:
116
116
  # for `CompatArray` input
117
117
  return arr.copy() if copy else arr
118
118
 
119
+ _check = kwargs.get("check", True)
119
120
  # for non-`CompatArray` input
120
- if not api.is_array_api_obj(arr):
121
+ if _check and not api.is_array_api_obj(arr):
121
122
  raise NotArrayAPIObjectError(
122
123
  f"Parameter `arr` of `CompatArray` must be an array API compatible array object, got {type(arr)}."
123
124
  )
@@ -135,6 +136,8 @@ class CompatArray(Compat):
135
136
  Convert `self` to a `NumPy array`.
136
137
  See also :func:`convert.to_numpy`.
137
138
  """
139
+ if self._xp_name == "NumPy" and not copy:
140
+ return self._arr
138
141
  return to_numpy(self._arr, copy=copy)
139
142
 
140
143
  def to_tensor(self, *, device=None, copy=False):
@@ -142,6 +145,8 @@ class CompatArray(Compat):
142
145
  Convert `self` to a `PyTorch tensor`.
143
146
  See also :func:`convert.to_tensor`.
144
147
  """
148
+ if self._xp_name == "PyTorch" and not copy and (device is None or device == self.device):
149
+ return self._arr
145
150
  return to_tensor(self._arr, device=device, copy=copy)
146
151
 
147
152
  def to_list(self, *, copy=False):
@@ -197,7 +202,7 @@ class CompatArray(Compat):
197
202
  All the arrays have the same shape.
198
203
  """
199
204
  result = self._get_xp_attr("unstack")(self._arr, axis=axis)
200
- return tuple(CompatArray(arr, xp=self._xp) for arr in result)
205
+ return tuple([CompatArray(arr, xp=self._cxp, check=False) for arr in result])
201
206
 
202
207
  # === Searching functions ===
203
208
  def nonzero(self):
@@ -218,7 +223,7 @@ class CompatArray(Compat):
218
223
  - If `self` has a boolean data type, non-zero elements are those elements which are equal to `True`.
219
224
  """
220
225
  result = self._get_xp_attr("nonzero")(self._arr)
221
- return tuple(CompatArray(arr, xp=self._xp) for arr in result)
226
+ return tuple([CompatArray(arr, xp=self._cxp) for arr in result])
222
227
 
223
228
  # === Set functions ===
224
229
  def unique_all(self):
@@ -238,10 +243,10 @@ class CompatArray(Compat):
238
243
  """
239
244
  result = self._get_xp_attr("unique_all")(self._arr)
240
245
  return UniqueResult(
241
- values=CompatArray(result.values, xp=self._xp),
242
- indices=CompatArray(result.indices, xp=self._xp),
243
- inverse_indices=CompatArray(result.inverse_indices, xp=self._xp),
244
- counts=CompatArray(result.counts, xp=self._xp),
246
+ values=CompatArray(result.values, xp=self._cxp, check=False),
247
+ indices=CompatArray(result.indices, xp=self._cxp, check=False),
248
+ inverse_indices=CompatArray(result.inverse_indices, xp=self._cxp, check=False),
249
+ counts=CompatArray(result.counts, xp=self._cxp, check=False),
245
250
  )
246
251
 
247
252
  def unique_counts(self):
@@ -259,10 +264,10 @@ class CompatArray(Compat):
259
264
  """
260
265
  result = self._get_xp_attr("unique_counts")(self._arr)
261
266
  return UniqueResult(
262
- values=CompatArray(result.values, xp=self._xp),
267
+ values=CompatArray(result.values, xp=self._cxp, check=False),
263
268
  indices=None,
264
269
  inverse_indices=None,
265
- counts=CompatArray(result.counts, xp=self._xp),
270
+ counts=CompatArray(result.counts, xp=self._cxp, check=False),
266
271
  )
267
272
 
268
273
  def unique_inverse(self):
@@ -280,9 +285,9 @@ class CompatArray(Compat):
280
285
  """
281
286
  result = self._get_xp_attr("unique_inverse")(self._arr)
282
287
  return UniqueResult(
283
- values=CompatArray(result.values, xp=self._xp),
288
+ values=CompatArray(result.values, xp=self._cxp, check=False),
284
289
  indices=None,
285
- inverse_indices=CompatArray(result.inverse_indices, xp=self._xp),
290
+ inverse_indices=CompatArray(result.inverse_indices, xp=self._cxp, check=False),
286
291
  counts=None,
287
292
  )
288
293
 
@@ -291,7 +296,7 @@ class CompatArray(Compat):
291
296
  """
292
297
  Return a copy of `self` via :func:`convert.as_array`.
293
298
  """
294
- return CompatArray.from_other(self._arr, xp=self._xp, copy=True)
299
+ return CompatArray.from_other(self._arr, xp=self._cxp, copy=True)
295
300
 
296
301
  def _get_attr(self, name: str):
297
302
  """Try to get the attribute `name` from `self`."""
@@ -334,7 +339,7 @@ class CompatArray(Compat):
334
339
  @property
335
340
  def device(self):
336
341
  """
337
- DeviceT on which `self` is stored.
342
+ Device on which `self` is stored.
338
343
  """
339
344
  return api.device(self._arr)
340
345
 
@@ -379,7 +384,7 @@ class CompatArray(Compat):
379
384
  result = self._get_xp_attr("T")(self._arr)
380
385
  except (AttributeError, TypeError):
381
386
  result = self._get_attr("T")
382
- return CompatArray(result, xp=self._xp)
387
+ return CompatArray(result, xp=self._cxp, check=False)
383
388
 
384
389
  @property
385
390
  def mT(self):
@@ -391,19 +396,19 @@ class CompatArray(Compat):
391
396
  result = self._get_xp_attr("mT")(self._arr)
392
397
  except (AttributeError, TypeError):
393
398
  result = self._get_attr("mT")
394
- return CompatArray(result, xp=self._xp)
399
+ return CompatArray(result, xp=self._cxp, check=False)
395
400
 
396
401
  def __array__(self):
397
402
  """Allow implicit NumPy conversion."""
398
403
  return self.to_numpy()
399
404
 
400
- def __getattr__(self, name: str):
401
- attr = self._get_cxp_attr(name)
405
+ def __getattr__(self, name):
406
+ attr = self._get_xp_attr(name)
402
407
 
403
408
  if callable(attr) and not isinstance(attr, type):
404
- def wrapper(*args, **kwargs):
405
- return attr(self._arr, *args, **kwargs)
406
- return wrapper
409
+ wrapped = _make_wrapper(self._xp_name, attr, self._cxp, self._arr)
410
+ self.__dict__[name] = wrapped
411
+ return wrapped
407
412
  raise CompatArrayAttributeError(f"`CompatArray` `{self._xp_name}` does not support attribute `{name}`.")
408
413
 
409
414
  def __len__(self):
@@ -540,7 +545,7 @@ def unwrap(obj):
540
545
  """
541
546
  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.
542
547
  """
543
- return obj.arr if isinstance(obj, CompatArray) else obj
548
+ return obj.arr if type(obj) is CompatArray else obj
544
549
 
545
550
 
546
551
  def wrap_arraylike(arr, xp=None):
@@ -549,8 +554,8 @@ def wrap_arraylike(arr, xp=None):
549
554
  """
550
555
  if api.is_array_api_obj(arr):
551
556
  if xp is None:
552
- return CompatArray(arr)
553
- return CompatArray(arr, xp=xp)
557
+ return CompatArray(arr, check=False)
558
+ return CompatArray(arr, xp=xp, check=False)
554
559
  return arr
555
560
 
556
561
 
@@ -559,3 +564,26 @@ def to_cxp(xp):
559
564
  from ._namespace import CompatNamespace
560
565
 
561
566
  return CompatNamespace(xp)
567
+
568
+
569
+ def _make_wrapper(xp_name, attr, cxp, first, /):
570
+ """Make a wrapper function for the attribute `name` of the `array namespace`."""
571
+ wrap_ = lambda x: wrap_arraylike(x, xp=cxp)
572
+ unwrap_ = unwrap
573
+
574
+ def wrapper(*args, **kwargs):
575
+ if xp_name == "NumPy":
576
+ return wrap_(attr(first, *args, **kwargs))
577
+
578
+ n = len(args)
579
+ if n == 0 and not kwargs:
580
+ return wrap_(attr(first))
581
+ if n == 1 and not kwargs:
582
+ return wrap_(attr(first, unwrap_(args[0])))
583
+ new_args = [unwrap_(a) for a in args]
584
+ if kwargs:
585
+ new_kwargs = {k: unwrap_(v) for k, v in kwargs.items()}
586
+ return wrap_(attr(first, *new_args, **new_kwargs))
587
+ return wrap_(attr(first, *new_args))
588
+ wrapper.__name__ = getattr(attr, "__name__", "wrapper")
589
+ return wrapper
@@ -10,7 +10,7 @@ from typing import (Union, List, Tuple, Optional, Any, Sequence, Generic, TypeVa
10
10
  from ._base import Compat
11
11
  from ._namespace import CompatNamespace
12
12
  from ..types import (
13
- T, DTypeT, DeviceT, dtypeT, deviceT, DType, AnyDevice,
13
+ T, dtypeT, DTypeT, deviceT, DeviceT, DType, AnyDevice,
14
14
  ArrayLike, ArrayLibraryName,
15
15
  ArrayOrAny, ArrayOrScalar, ArrayOrReal, ArrayOrIntLike, ArrayOrInt, ArrayOrbool,
16
16
  UniqueAllResult, UniqueCountsResult, UniqueInverseResult
@@ -64,7 +64,6 @@ class CompatArray(Compat, Generic[TT, DT]):
64
64
  def to_device(self, device: DeviceT, /, *, stream: Optional[Any] = None) -> CompatArray[TT, DeviceT]: ...
65
65
  @overload
66
66
  def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[TT, AnyDevice]: ...
67
-
68
67
  def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[Any, AnyDevice]: ...
69
68
 
70
69
  # === Data type functions ===
@@ -1771,6 +1770,8 @@ class CompatArray(Compat, Generic[TT, DT]):
1771
1770
 
1772
1771
  def __array__(self) -> NDArray[TT]: ...
1773
1772
 
1773
+ def __getattr__(self, name: str) -> Any: ...
1774
+
1774
1775
  def __len__(self) -> int: ...
1775
1776
  def __abs__(self) -> CompatArray[TT, DT]: ...
1776
1777
  def __add__(self, other: ArrayOrScalar, /) -> CompatArray[Any, DT]: ...
@@ -3,7 +3,7 @@
3
3
  # @TianZhen
4
4
 
5
5
  from __future__ import annotations
6
- from typing import (Any, TYPE_CHECKING)
6
+ from typing import (Any, TYPE_CHECKING, Callable)
7
7
 
8
8
  from .._utils import array_namespace_alias
9
9
 
@@ -50,7 +50,7 @@ class CompatNamespace(Compat):
50
50
  AttributeError: ...
51
51
  """
52
52
  def __new__(cls, xp, /):
53
- if isinstance(xp, CompatNamespace):
53
+ if type(xp) is cls:
54
54
  # for `CompatNamespace` input
55
55
  return xp
56
56
  # for `Namespace` input
@@ -91,10 +91,10 @@ class CompatNamespace(Compat):
91
91
  Each returned array should have the same data type as the input arrays.
92
92
  """
93
93
  result = self._get_xp_attr("meshgrid")(
94
- *tuple(unwrap(arr) for arr in arrays),
94
+ *[unwrap(arr) for arr in arrays],
95
95
  indexing=indexing
96
96
  )
97
- return [CompatArray(arr, xp=self._xp) for arr in result]
97
+ return [CompatArray(arr, xp=self) for arr in result]
98
98
 
99
99
  # === Data Type functions ===
100
100
  def can_cast(self, from_, to, /):
@@ -170,7 +170,7 @@ class CompatNamespace(Compat):
170
170
  """
171
171
  return self._get_xp_attr("iinfo")(unwrap(type_))
172
172
 
173
- def isdtype(self, dtype, kind) -> bool:
173
+ def isdtype(self, dtype, kind):
174
174
  """
175
175
  Returns a boolean indicating whether a provided :param:`dtype` is of a specified data type :param:`kind`.
176
176
 
@@ -214,7 +214,7 @@ class CompatNamespace(Compat):
214
214
  DType
215
215
  The dtype resulting from an operation involving the input arrays, scalars, and/or dtypes.
216
216
  """
217
- return self._get_xp_attr("result_type")(*tuple(unwrap(arr) for arr in arrays_and_dtypes))
217
+ return self._get_xp_attr("result_type")(*[unwrap(arr) for arr in arrays_and_dtypes])
218
218
 
219
219
  # === Manipulation functions ===
220
220
  def broadcast_arrays(self, *arrays):
@@ -233,8 +233,139 @@ class CompatNamespace(Compat):
233
233
  Each array must have the same shape.
234
234
  Each array must have the same dtype as its corresponding input array.
235
235
  """
236
- result = self._get_xp_attr("broadcast_arrays")(*tuple(unwrap(arr) for arr in arrays))
237
- return [CompatArray(arr, xp=self._xp) for arr in result]
236
+ result = self._get_xp_attr("broadcast_arrays")(*[unwrap(arr) for arr in arrays])
237
+ return [CompatArray(arr, xp=self) for arr in result]
238
+
239
+ # === Linear Algebra Extension ===
240
+ def vector_norm(self, x, /, *, axis=None, keepdims=False, ord=2):
241
+ """
242
+ Computes the vector norm of a vector (or batch of vectors) :param:`x`.
243
+
244
+ Parameters
245
+ ----------
246
+ x : ArrayLike[Any]
247
+ The input array. Should have a floating-point data type.
248
+
249
+ axis : Optional[Union[int, Tuple[int, ...]]], default to `None`
250
+ - _int_: :param:`axis` specifies the axis (dimension) along which to compute vector norms;
251
+ - _tuple_: :param:`axis` specifies the axes (dimensions) along which to compute batched vector norms;
252
+ - `None`: The vector norm must be computed over all array values (i.e., equivalent to computing the vector norm of a flattened array).
253
+
254
+ Negative indices must be supported.
255
+
256
+ keepdims : bool, default to `False`
257
+ - `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);
258
+ - `False`: The axes (dimensions) specified by :param:`axis` must not be included in the result.
259
+
260
+ ord : Union[int, float, Literal['inf', '-inf']], default to `2`
261
+ Order of the norm.
262
+ The following mathematical norms must be supported:
263
+ - `1`: L1-norm (Manhattan);
264
+ - `2`: L2-norm (Euclidean);
265
+ - `"inf"`: infinity norm;
266
+ - _int_ or _float_ (>=1): p-norm.
267
+
268
+ The following non-mathematical “norms” must be supported:
269
+ - `0`: sum(a != 0);
270
+ - `-1`: 1./sum(1./abs(a));
271
+ - `-2`: 1./sqrt(sum(1./a**2));
272
+ - `"-inf"`: min(abs(a));
273
+ - _int_ or _float_ (<1): sum(abs(a)**ord)**(1./ord).
274
+
275
+ Returns
276
+ -------
277
+ CompatArray
278
+ A :class:`CompatArray` array containing the vector norms.
279
+ - :param:`axis` is `None`: The returned array must be a zero-dimensional array containing a vector norm;
280
+ - :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`;
281
+ - :param:`axis` is _tuple_ (`n` elements): The returned array must have a rank which is `n` less than the rank of :param:`x`;
282
+
283
+ - :param:`x` is real-valued data type: The returned array must have a real-valued floating-point data type determined by Type Promotion Rules;
284
+ - :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).
285
+ """
286
+ if ord == "inf":
287
+ ord = float("inf")
288
+ elif ord == "-inf":
289
+ ord = float("-inf")
290
+
291
+ result = getattr(self.linalg, "vector_norm")(unwrap(x), axis=axis, keepdims=keepdims, ord=ord)
292
+ return CompatArray(result, xp=self)
293
+
294
+ def matrix_norm(self, x, /, *, keepdims=False, ord="fro"):
295
+ """
296
+ Computes the matrix norm of a matrix (or a stack of matrices) :param:`x`.
297
+
298
+ Parameters
299
+ ----------
300
+ x : ArrayLike[Any]
301
+ Input array having shape (..., `M`, `N`) and whose innermost two dimensions form `MxN` matrices. Should have a floating-point data type.
302
+
303
+ keepdims : bool, default to `False`
304
+ - `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);
305
+ - `False`: The last two axes (dimensions) must not be included in the result.
306
+
307
+ ord : Optional[Union[int, float, Literal['inf', '-inf', 'fro', 'nuc']]], default to `"fro"`
308
+ order of the norm.
309
+ The following mathematical norms must be supported:
310
+ - `"fro"`: Frobenius norm;
311
+ - `"nuc"`: nuclear norm;
312
+ - `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);
313
+ - `2`: largest singular value. The norm corresponds to the induced matrix norm where `p=inf` (i.e., the maximum absolute value row sum);
314
+ - `"inf"`: max(sum(abs(x), axis=1)). The norm corresponds to the induced matrix norm where `p=2` (i.e., the largest singular value).
315
+
316
+ The following non-mathematical “norms” must be supported:
317
+ - `-1`: min(sum(abs(x), axis=0));
318
+ - `-2`: smallest singular value;
319
+ - `"-inf"`: min(sum(abs(x), axis=1)).
320
+
321
+ Returns
322
+ -------
323
+ CompatArray
324
+ A :class:`CompatArray` array containing the norms for each `MxN` matrix.
325
+ - :param:`keepdims` is `False`: The returned array must have a rank which is two less than the rank of :param:`x`;
326
+
327
+ - :param:`x` is real-valued data type: The returned array must have a real-valued floating-point data type determined by Type Promotion Rules;
328
+ - :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).
329
+
330
+ """
331
+ if ord == "inf":
332
+ ord = float("inf")
333
+ elif ord == "-inf":
334
+ ord = float("-inf")
335
+
336
+ result = getattr(self.linalg, "matrix_norm")(unwrap(x), keepdims=keepdims, ord=ord)
337
+ return CompatArray(result, xp=self)
338
+
339
+ @property
340
+ def linalg(self):
341
+ """
342
+ The `linalg` namespace for linear algebra functions.
343
+ The following functions must be supported in the `linalg` namespace:
344
+ - `cholesky`(x, /, *, upper=False): Returns the lower (upper) Cholesky decomposition of a complex Hermitian or real symmetric positive-definite matrix x.
345
+ - `cross`(x1, x2, /, *, axis=-1): Returns the cross product of 3-element vectors.
346
+ - `det`(x, /): Returns the determinant of a square matrix (or a stack of square matrices) x.
347
+ - `diagonal`(x, /, *, offset=0): Returns the specified diagonals of a matrix (or a stack of matrices) x.
348
+ - `eigh`(x, /): Returns an eigenvalue decomposition of a complex Hermitian or real symmetric matrix (or a stack of matrices) x.
349
+ - `eigvalsh`(x, /): Returns the eigenvalues of a complex Hermitian or real symmetric matrix (or a stack of matrices) x.
350
+ - `inv`(x, /): Returns the multiplicative inverse of a square matrix (or a stack of square matrices) x.
351
+ - `matmul`(x1, x2, /): Alias for matmul().
352
+ - `matrix_norm`(x, /, *, keepdims=False, ord='fro'): Computes the matrix norm of a matrix (or a stack of matrices) x.
353
+ - `matrix_power`(x, n, /): Raises a square matrix (or a stack of square matrices) x to an integer power n.
354
+ - `matrix_rank`(x, /, *, rtol=None): Returns the rank (i.e., number of non-zero singular values) of a matrix (or a stack of matrices).
355
+ - `matrix_transpose`(x, /): Alias for matrix_transpose().
356
+ - `outer`(x1, x2, /): Returns the outer product of two vectors x1 and x2.
357
+ - `pinv`(x, /, *, rtol=None): Returns the (Moore-Penrose) pseudo-inverse of a matrix (or a stack of matrices) x.
358
+ - `qr`(x, /, *, mode='reduced'): Returns the QR decomposition of a full column rank matrix (or a stack of matrices).
359
+ - `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.
360
+ - `solve`(x1, x2, /): Returns the solution of a square system of linear equations with a unique solution.
361
+ - `svd`(x, /, *, full_matrices=True): Returns a singular value decomposition (SVD) of a matrix (or a stack of matrices) x.
362
+ - `svdvals`(x, /): Returns the singular values of a matrix (or a stack of matrices) x.
363
+ - `tensordot`(x1, x2, /, *, axes=2): Alias for tensordot().
364
+ - `trace`(x, /, *, offset=0, dtype=None): Returns the sum along the specified diagonals of a matrix (or a stack of matrices) x.
365
+ - `vecdot`(x1, x2, /, *, axis=-1): Alias for vecdot().
366
+ - `vector_norm`(x, /, *, axis=None, keepdims=False, ord=2): Computes the vector norm of a vector (or batch of vectors) x.
367
+ """
368
+ return self._get_xp_attr("linalg")
238
369
 
239
370
  # === Constants ===
240
371
  @property
@@ -315,16 +446,37 @@ class CompatNamespace(Compat):
315
446
  def __name__(self):
316
447
  return "(compat)" + getattr(self._xp, "__name__", type(self._xp).__name__)
317
448
 
318
- def __getattr__(self, name: str):
449
+ def __getattr__(self, name):
319
450
  attr = self._get_xp_attr(name)
320
451
 
321
452
  if callable(attr):
322
- def wrapper(*args, **kwargs):
323
- if not args and not kwargs:
324
- return wrap_arraylike(attr(), xp=self._xp)
325
-
326
- new_args = tuple(unwrap(a) for a in args)
327
- new_kwargs = {k: unwrap(v) for k, v in kwargs.items()} if kwargs else kwargs
328
- return wrap_arraylike(attr(*new_args, **new_kwargs), xp=self._xp)
329
- return wrapper
453
+ wrapped = _make_wrapper(self._xp_name, attr, self)
454
+ self.__dict__[name] = wrapped
455
+ return wrapped
330
456
  raise CompatNamespaceAttributeError(f"`CompatNamespace` `{self._xp_name}` does not support attribute `{name}`.")
457
+
458
+
459
+ def _make_wrapper(xp_name, attr, cxp, /):
460
+ """Make a wrapper function for the attribute `name` of the `array namespace`."""
461
+ wrap_ = lambda x: wrap_arraylike(x, xp=cxp)
462
+ unwrap_ = unwrap
463
+
464
+ def wrapper(*args, **kwargs):
465
+ if xp_name == "NumPy":
466
+ return wrap_(attr(*args, **kwargs))
467
+
468
+ n = len(args)
469
+ if n == 0 and not kwargs:
470
+ return wrap_(attr())
471
+ if n == 1 and not kwargs:
472
+ return wrap_(attr(unwrap_(args[0])))
473
+ if n == 2 and not kwargs:
474
+ a0, a1 = args
475
+ return wrap_(attr(unwrap_(a0), unwrap_(a1)))
476
+ new_args = [unwrap_(a) for a in args]
477
+ if kwargs:
478
+ new_kwargs = {k: unwrap_(v) for k, v in kwargs.items()}
479
+ return wrap_(attr(*new_args, **new_kwargs))
480
+ return wrap_(attr(*new_args))
481
+ wrapper.__name__ = getattr(attr, "__name__", "wrapper")
482
+ return wrapper
@@ -4,12 +4,13 @@
4
4
 
5
5
  from __future__ import annotations
6
6
  from numpy.typing import NDArray
7
+ from types import ModuleType
7
8
  from typing import (Union, List, Tuple, Optional, Any, Literal, overload)
8
9
 
9
10
  from ._base import Compat
10
11
  from ._array import CompatArray
11
12
  from ..types import (
12
- DTypeT, DeviceT, dtypeT, DType, AnyDevice,
13
+ dtypeT, DTypeT, deviceT, DeviceT, DType, AnyDevice,
13
14
  ValueT, Value, ArrayLike, ArrayOrAny
14
15
  )
15
16
 
@@ -802,7 +803,7 @@ class CompatNamespace(Compat):
802
803
  Returns
803
804
  -------
804
805
  CompatArray
805
- A :class:`CompatArray` output array containing the concatenated values.
806
+ A :class:`CompatArray` output array containing the concatenated values.
806
807
  """
807
808
  ...
808
809
 
@@ -841,6 +842,26 @@ class CompatNamespace(Compat):
841
842
  """
842
843
  ...
843
844
 
845
+ # === Linear Algebra Extension ===
846
+ @overload
847
+ 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"]]: ...
848
+ @overload
849
+ 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]: ...
850
+ @overload
851
+ 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]: ...
852
+ 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]: ...
853
+
854
+ @overload
855
+ def matrix_norm(self, x: NDArray[Any], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, Literal["cpu"]]: ...
856
+ @overload
857
+ def matrix_norm(self, x: CompatArray[Any, deviceT], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, deviceT]: ...
858
+ @overload
859
+ def matrix_norm(self, x: ArrayLike[Any], /, *, keepdims: bool = ..., ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = ...) -> CompatArray[float, AnyDevice]: ...
860
+ def matrix_norm(self, x: ArrayLike[Any], /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal["inf", "-inf", "fro", "nuc"]]] = "fro") -> CompatArray[float, AnyDevice]: ...
861
+
862
+ @property
863
+ def linalg(self) -> ModuleType: ...
864
+
844
865
  # === Constants ===
845
866
  @property
846
867
  def e(self) -> float: ...
@@ -883,3 +904,5 @@ class CompatNamespace(Compat):
883
904
 
884
905
  @property
885
906
  def __name__(self) -> str: ...
907
+
908
+ def __getattr__(self, name: str) -> Any: ...
@@ -122,11 +122,11 @@ def to_tensor(obj, /, *, dtype=None, device=None, copy=True):
122
122
  - `None`: Raises `ConvertNoneTypeError`;
123
123
  - _others_: Converted to a `PyTorch tensor` directly.
124
124
 
125
- dtype : Optional[DTypeT], default to `None`
125
+ dtype : Optional[DType], default to `None`
126
126
  The data type of the resulting `PyTorch tensor`.
127
127
  - `None`: Use the default data type of the object.
128
128
 
129
- device : Optional[DeviceT], default to `None`
129
+ device : Optional[AnyDevice], default to `None`
130
130
  The device on which the resulting `PyTorch tensor` will be allocated.
131
131
  - `None`: Use the default device (usually `"cpu"`).
132
132
 
@@ -294,11 +294,11 @@ def as_array(obj, xp, /, *, dtype=None, device=None, copy=False, arraylike_only=
294
294
  - _ArrayLibraryName_ (`"numpy"` or `"torch"`): Converted to a `NumPy array` or `PyTorch tensor` respectively using the corresponding conversion functions;
295
295
  - _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.
296
296
 
297
- dtype : Optional[DTypeT], default to `None`
297
+ dtype : Optional[DType], default to `None`
298
298
  The data type of the resulting array.
299
299
  - `None`: Use the default data type of the object.
300
300
 
301
- device : Optional[DeviceT], default to `None`
301
+ device : Optional[AnyDevice], default to `None`
302
302
  The device on which the resulting array will be allocated (only if `array namespace` supports it).
303
303
 
304
304
  copy : bool, default to `False`