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.
- {cobra_array-0.1.5/src/cobra_array.egg-info → cobra_array-0.2.1}/PKG-INFO +1 -1
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/__init__.py +1 -1
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/_core.py +2 -2
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/_utils.py +27 -12
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_array.py +53 -25
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_array.pyi +3 -2
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_base.py +1 -1
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_namespace.py +168 -16
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/_namespace.pyi +25 -2
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/convert.py +4 -4
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/types.py +3 -2
- {cobra_array-0.1.5 → cobra_array-0.2.1/src/cobra_array.egg-info}/PKG-INFO +1 -1
- cobra_array-0.2.1/tests/test_compat.py +523 -0
- cobra_array-0.1.5/tests/test_compat.py +0 -261
- {cobra_array-0.1.5 → cobra_array-0.2.1}/LICENSE +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/README.md +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/pyproject.toml +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/setup.cfg +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/array_api.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/compat/__init__.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/convert.pyi +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/default.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array/exceptions.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/SOURCES.txt +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/dependency_links.txt +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/requires.txt +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/src/cobra_array.egg-info/top_level.txt +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_backend.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_compat_namespace.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_convert.py +0 -0
- {cobra_array-0.1.5 → cobra_array-0.2.1}/tests/test_default.py +0 -0
- {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
|
|
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
|
|
@@ -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 =
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
74
|
-
|
|
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
|
-
|
|
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
|
|
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.
|
|
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.
|
|
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.
|
|
242
|
-
indices=CompatArray(result.indices, xp=self.
|
|
243
|
-
inverse_indices=CompatArray(result.inverse_indices, xp=self.
|
|
244
|
-
counts=CompatArray(result.counts, xp=self.
|
|
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.
|
|
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.
|
|
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.
|
|
288
|
+
values=CompatArray(result.values, xp=self._cxp, check=False),
|
|
284
289
|
indices=None,
|
|
285
|
-
inverse_indices=CompatArray(result.inverse_indices, xp=self.
|
|
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.
|
|
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
|
-
|
|
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.
|
|
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.
|
|
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
|
|
401
|
-
attr = self.
|
|
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
|
-
|
|
405
|
-
|
|
406
|
-
return
|
|
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
|
|
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,
|
|
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]: ...
|
|
@@ -50,7 +50,7 @@ class CompatNamespace(Compat):
|
|
|
50
50
|
AttributeError: ...
|
|
51
51
|
"""
|
|
52
52
|
def __new__(cls, xp, /):
|
|
53
|
-
if
|
|
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
|
-
*
|
|
94
|
+
*[unwrap(arr) for arr in arrays],
|
|
95
95
|
indexing=indexing
|
|
96
96
|
)
|
|
97
|
-
return [CompatArray(arr, xp=self
|
|
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)
|
|
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")(*
|
|
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")(*
|
|
237
|
-
return [CompatArray(arr, xp=self
|
|
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
|
|
449
|
+
def __getattr__(self, name):
|
|
319
450
|
attr = self._get_xp_attr(name)
|
|
320
451
|
|
|
321
452
|
if callable(attr):
|
|
322
|
-
|
|
323
|
-
|
|
324
|
-
|
|
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,
|
|
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[
|
|
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[
|
|
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[
|
|
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[
|
|
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`
|