cobra-array 0.1.0__tar.gz → 0.1.2__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.0/src/cobra_array.egg-info → cobra_array-0.1.2}/PKG-INFO +3 -11
- {cobra_array-0.1.0 → cobra_array-0.1.2}/README.md +2 -10
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/__init__.py +7 -15
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/_core.py +84 -10
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/_utils.py +1 -1
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/array_api.py +67 -1
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/__init__.py +17 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_array.py +56 -1
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_array.pyi +24 -17
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_namespace.py +25 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_namespace.pyi +122 -70
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/convert.py +53 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/convert.pyi +13 -13
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/default.py +30 -11
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/types.py +17 -15
- {cobra_array-0.1.0 → cobra_array-0.1.2/src/cobra_array.egg-info}/PKG-INFO +3 -11
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/SOURCES.txt +2 -0
- cobra_array-0.1.2/tests/test_compat.py +253 -0
- cobra_array-0.1.2/tests/test_compat_namespace.py +122 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_default.py +1 -1
- {cobra_array-0.1.0 → cobra_array-0.1.2}/LICENSE +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/pyproject.toml +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/setup.cfg +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_base.py +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/exceptions.py +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/dependency_links.txt +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/requires.txt +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/top_level.txt +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_backend.py +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_convert.py +0 -0
- {cobra_array-0.1.0 → cobra_array-0.1.2}/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.1.2
|
|
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
|
|
@@ -71,14 +71,11 @@ pip install cobra-array
|
|
|
71
71
|
|
|
72
72
|
data = [[1, 2], [3, 4]]
|
|
73
73
|
|
|
74
|
-
arr_np = to_numpy(data, dtype=np.float32)
|
|
75
|
-
print(type(arr_np), arr_np.dtype) # numpy.ndarray float32
|
|
74
|
+
arr_np = to_numpy(data, dtype=np.float32) # numpy.ndarray float32
|
|
76
75
|
|
|
77
76
|
arr_torch = to_tensor(data, device="cpu")
|
|
78
|
-
print(type(arr_torch), arr_torch.device)
|
|
79
77
|
|
|
80
|
-
back_to_list = to_list(arr_np)
|
|
81
|
-
print(back_to_list) # [[1.0, 2.0], [3.0, 4.0]]
|
|
78
|
+
back_to_list = to_list(arr_np) # [[1.0, 2.0], [3.0, 4.0]]
|
|
82
79
|
```
|
|
83
80
|
|
|
84
81
|
- Context-based conversion:
|
|
@@ -91,8 +88,6 @@ pip install cobra-array
|
|
|
91
88
|
x = as_context([1, 2, 3])
|
|
92
89
|
y = as_context(np.array([4, 5]))
|
|
93
90
|
spec = context_spec()
|
|
94
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
95
|
-
print(x, y)
|
|
96
91
|
```
|
|
97
92
|
|
|
98
93
|
- Auto-unify function arguments:
|
|
@@ -107,7 +102,6 @@ pip install cobra-array
|
|
|
107
102
|
return c.mean()
|
|
108
103
|
|
|
109
104
|
out = add_and_mean(np.array([1, 2, 3]), [4, 5, 6])
|
|
110
|
-
print(out)
|
|
111
105
|
```
|
|
112
106
|
|
|
113
107
|
- Default backend strategy:
|
|
@@ -116,10 +110,8 @@ pip install cobra-array
|
|
|
116
110
|
from cobra_array.default import as_default, default_spec
|
|
117
111
|
|
|
118
112
|
spec = default_spec()
|
|
119
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
120
113
|
|
|
121
114
|
x = as_default([1, 2, 3], unify_dtype=True, unify_device=True)
|
|
122
|
-
print(x, x.dtype)
|
|
123
115
|
```
|
|
124
116
|
|
|
125
117
|
## Requirements
|
|
@@ -46,14 +46,11 @@ pip install cobra-array
|
|
|
46
46
|
|
|
47
47
|
data = [[1, 2], [3, 4]]
|
|
48
48
|
|
|
49
|
-
arr_np = to_numpy(data, dtype=np.float32)
|
|
50
|
-
print(type(arr_np), arr_np.dtype) # numpy.ndarray float32
|
|
49
|
+
arr_np = to_numpy(data, dtype=np.float32) # numpy.ndarray float32
|
|
51
50
|
|
|
52
51
|
arr_torch = to_tensor(data, device="cpu")
|
|
53
|
-
print(type(arr_torch), arr_torch.device)
|
|
54
52
|
|
|
55
|
-
back_to_list = to_list(arr_np)
|
|
56
|
-
print(back_to_list) # [[1.0, 2.0], [3.0, 4.0]]
|
|
53
|
+
back_to_list = to_list(arr_np) # [[1.0, 2.0], [3.0, 4.0]]
|
|
57
54
|
```
|
|
58
55
|
|
|
59
56
|
- Context-based conversion:
|
|
@@ -66,8 +63,6 @@ pip install cobra-array
|
|
|
66
63
|
x = as_context([1, 2, 3])
|
|
67
64
|
y = as_context(np.array([4, 5]))
|
|
68
65
|
spec = context_spec()
|
|
69
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
70
|
-
print(x, y)
|
|
71
66
|
```
|
|
72
67
|
|
|
73
68
|
- Auto-unify function arguments:
|
|
@@ -82,7 +77,6 @@ pip install cobra-array
|
|
|
82
77
|
return c.mean()
|
|
83
78
|
|
|
84
79
|
out = add_and_mean(np.array([1, 2, 3]), [4, 5, 6])
|
|
85
|
-
print(out)
|
|
86
80
|
```
|
|
87
81
|
|
|
88
82
|
- Default backend strategy:
|
|
@@ -91,10 +85,8 @@ pip install cobra-array
|
|
|
91
85
|
from cobra_array.default import as_default, default_spec
|
|
92
86
|
|
|
93
87
|
spec = default_spec()
|
|
94
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
95
88
|
|
|
96
89
|
x = as_default([1, 2, 3], unify_dtype=True, unify_device=True)
|
|
97
|
-
print(x, x.dtype)
|
|
98
90
|
```
|
|
99
91
|
|
|
100
92
|
## Requirements
|
|
@@ -25,23 +25,20 @@ Functions
|
|
|
25
25
|
|
|
26
26
|
Examples
|
|
27
27
|
--------
|
|
28
|
-
Basic conversions::
|
|
28
|
+
- Basic conversions::
|
|
29
29
|
|
|
30
30
|
import numpy as np
|
|
31
31
|
from cobra_array.convert import to_numpy, to_tensor, to_list
|
|
32
32
|
|
|
33
33
|
data = [[1, 2], [3, 4]]
|
|
34
34
|
|
|
35
|
-
arr_np = to_numpy(data, dtype=np.float32)
|
|
36
|
-
print(type(arr_np), arr_np.dtype) # numpy.ndarray float32
|
|
35
|
+
arr_np = to_numpy(data, dtype=np.float32) # numpy.ndarray float32
|
|
37
36
|
|
|
38
37
|
arr_torch = to_tensor(data, device="cpu")
|
|
39
|
-
print(type(arr_torch), arr_torch.device)
|
|
40
38
|
|
|
41
|
-
back_to_list = to_list(arr_np)
|
|
42
|
-
print(back_to_list) # [[1.0, 2.0], [3.0, 4.0]]
|
|
39
|
+
back_to_list = to_list(arr_np) # [[1.0, 2.0], [3.0, 4.0]]
|
|
43
40
|
|
|
44
|
-
Context-based conversion::
|
|
41
|
+
- Context-based conversion::
|
|
45
42
|
|
|
46
43
|
import numpy as np
|
|
47
44
|
from cobra_array import array_context, as_context, context_spec
|
|
@@ -50,10 +47,8 @@ Context-based conversion::
|
|
|
50
47
|
x = as_context([1, 2, 3])
|
|
51
48
|
y = as_context(np.array([4, 5]))
|
|
52
49
|
spec = context_spec()
|
|
53
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
54
|
-
print(x, y)
|
|
55
50
|
|
|
56
|
-
Auto-unify function arguments::
|
|
51
|
+
- Auto-unify function arguments::
|
|
57
52
|
|
|
58
53
|
import numpy as np
|
|
59
54
|
from cobra_array import unify_args
|
|
@@ -64,17 +59,14 @@ Auto-unify function arguments::
|
|
|
64
59
|
return c.mean()
|
|
65
60
|
|
|
66
61
|
out = add_and_mean(np.array([1, 2, 3]), [4, 5, 6])
|
|
67
|
-
print(out)
|
|
68
62
|
|
|
69
|
-
Default backend strategy::
|
|
63
|
+
- Default backend strategy::
|
|
70
64
|
|
|
71
65
|
from cobra_array.default import as_default, default_spec
|
|
72
66
|
|
|
73
67
|
spec = default_spec()
|
|
74
|
-
print(spec.cxp.xp_name, spec.dtype, spec.device)
|
|
75
68
|
|
|
76
69
|
x = as_default([1, 2, 3], unify_dtype=True, unify_device=True)
|
|
77
|
-
print(x, x.dtype)
|
|
78
70
|
"""
|
|
79
71
|
|
|
80
72
|
from ._core import (
|
|
@@ -91,7 +83,7 @@ from ._utils import (
|
|
|
91
83
|
)
|
|
92
84
|
|
|
93
85
|
__author__ = "Zhen Tian"
|
|
94
|
-
__version__ = "0.1.
|
|
86
|
+
__version__ = "0.1.2"
|
|
95
87
|
|
|
96
88
|
__all__ = [
|
|
97
89
|
"array_spec",
|
|
@@ -22,7 +22,7 @@ from .exceptions import (
|
|
|
22
22
|
if TYPE_CHECKING:
|
|
23
23
|
from numpy.typing import NDArray
|
|
24
24
|
from .compat import CompatArray
|
|
25
|
-
from .types import (T, dtypeT, DType,
|
|
25
|
+
from .types import (T, dtypeT, DType, AnyDevice, ArrayLike, ArrayLibraryName)
|
|
26
26
|
|
|
27
27
|
|
|
28
28
|
def array_spec(
|
|
@@ -88,7 +88,29 @@ def array_spec(
|
|
|
88
88
|
|
|
89
89
|
Examples
|
|
90
90
|
--------
|
|
91
|
-
|
|
91
|
+
Infer namespace from positional arrays:
|
|
92
|
+
|
|
93
|
+
>>> import numpy as np
|
|
94
|
+
>>> spec = array_spec(np.asarray([1, 2]), np.asarray([3, 4]))
|
|
95
|
+
>>> spec.cxp.xp_name
|
|
96
|
+
'NumPy'
|
|
97
|
+
>>> spec.dtype is None and spec.device is None
|
|
98
|
+
True
|
|
99
|
+
|
|
100
|
+
Use a keyword argument as reference to carry dtype and device:
|
|
101
|
+
|
|
102
|
+
>>> ref_arr = np.asarray([1, 2], dtype=np.float32)
|
|
103
|
+
>>> spec = array_spec(kw_arrays={"x": ref_arr}, ref="x")
|
|
104
|
+
>>> str(spec.dtype)
|
|
105
|
+
'float32'
|
|
106
|
+
>>> str(spec.device)
|
|
107
|
+
'cpu'
|
|
108
|
+
|
|
109
|
+
Filter non-array-like inputs when selecting reference by index:
|
|
110
|
+
|
|
111
|
+
>>> spec = array_spec("skip", np.asarray([1, 2]), ref=0, filter_arraylike=True)
|
|
112
|
+
>>> spec.cxp.xp_name
|
|
113
|
+
'NumPy'
|
|
92
114
|
"""
|
|
93
115
|
kw_arrays = kw_arrays or {}
|
|
94
116
|
|
|
@@ -197,11 +219,11 @@ def context_spec() -> ArraySpec:
|
|
|
197
219
|
@overload
|
|
198
220
|
def as_context(obj: NDArray[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT, Literal["cpu"]]: ...
|
|
199
221
|
@overload
|
|
200
|
-
def as_context(obj: ArrayLike[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT,
|
|
222
|
+
def as_context(obj: ArrayLike[dtypeT], /, *, unify_dtype: Literal[False], unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[dtypeT, AnyDevice]: ...
|
|
201
223
|
@overload
|
|
202
|
-
def as_context(obj: ArrayLike[Any], /, *, unify_dtype: Literal[True] = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[Any,
|
|
224
|
+
def as_context(obj: ArrayLike[Any], /, *, unify_dtype: Literal[True] = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: bool = ...) -> CompatArray[Any, AnyDevice]: ...
|
|
203
225
|
@overload
|
|
204
|
-
def as_context(obj: object, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[False] = ...) -> CompatArray[Any,
|
|
226
|
+
def as_context(obj: object, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[False] = ...) -> CompatArray[Any, AnyDevice]: ...
|
|
205
227
|
@overload
|
|
206
228
|
def as_context(obj: T, /, *, unify_dtype: bool = ..., unify_device: bool = ..., copy: bool = ..., arraylike_only: Literal[True]) -> T: ...
|
|
207
229
|
|
|
@@ -236,7 +258,7 @@ def as_context(
|
|
|
236
258
|
|
|
237
259
|
Returns
|
|
238
260
|
-------
|
|
239
|
-
CompatArray[Any,
|
|
261
|
+
CompatArray[Any, AnyDevice]
|
|
240
262
|
The converted array representation of the object in the current context `compatibility namespace`, with the current context `dtype` and `device` if specified.
|
|
241
263
|
object
|
|
242
264
|
If :param:`arraylike_only` is `True` and the object is not array-like.
|
|
@@ -244,6 +266,12 @@ def as_context(
|
|
|
244
266
|
Raises
|
|
245
267
|
------
|
|
246
268
|
Refer to :func:`convert.as_array`, :func:`context_spec` for possible exceptions.
|
|
269
|
+
|
|
270
|
+
Examples
|
|
271
|
+
--------
|
|
272
|
+
>>> from cobra_array import as_context
|
|
273
|
+
>>> as_context([1, 2, 3])
|
|
274
|
+
PyTorch_Array(tensor([1., 2., 3.], dtype=torch.float64))
|
|
247
275
|
"""
|
|
248
276
|
spec = context_spec()
|
|
249
277
|
return wrap_arraylike(as_array(
|
|
@@ -258,6 +286,40 @@ def as_context(
|
|
|
258
286
|
class array_context:
|
|
259
287
|
"""
|
|
260
288
|
**Context Manager** to set the context `compatibility namespace`, `dtype` and `device` for the enclosed block of code.
|
|
289
|
+
|
|
290
|
+
Examples
|
|
291
|
+
--------
|
|
292
|
+
Set a temporary context explicitly:
|
|
293
|
+
|
|
294
|
+
>>> from cobra_array import array_context
|
|
295
|
+
>>> with array_context(xp="numpy", dtype=None, device="cpu") as spec:
|
|
296
|
+
... spec.cxp.xp_name
|
|
297
|
+
'NumPy'
|
|
298
|
+
|
|
299
|
+
Convert data in the active context:
|
|
300
|
+
|
|
301
|
+
>>> with array_context(xp="numpy", dtype=None, device="cpu"):
|
|
302
|
+
... as_context([1, 2, 3], unify_dtype=False)
|
|
303
|
+
NumPy_Array([1 2 3])
|
|
304
|
+
|
|
305
|
+
Build a context from :class:`ArraySpec`:
|
|
306
|
+
|
|
307
|
+
>>> import numpy as np
|
|
308
|
+
>>> spec = array_spec(np.asarray([1, 2], dtype=np.float32), ref=0)
|
|
309
|
+
>>> with array_context.from_array_spec(spec) as cur:
|
|
310
|
+
... str(cur.dtype)
|
|
311
|
+
'float32'
|
|
312
|
+
|
|
313
|
+
Nested contexts override temporarily and then restore outer context:
|
|
314
|
+
|
|
315
|
+
>>> from cobra_array import context_spec
|
|
316
|
+
>>> with array_context(xp="numpy", dtype=None, device="cpu"):
|
|
317
|
+
... before = context_spec().dtype
|
|
318
|
+
... with array_context(dtype=float):
|
|
319
|
+
... middle = context_spec().dtype
|
|
320
|
+
... after = context_spec().dtype
|
|
321
|
+
... (before is None, middle is float, after is None)
|
|
322
|
+
(True, True, True)
|
|
261
323
|
"""
|
|
262
324
|
@classmethod
|
|
263
325
|
def from_array_spec(cls, arr_spec: ArraySpec, /):
|
|
@@ -275,7 +337,7 @@ class array_context:
|
|
|
275
337
|
self,
|
|
276
338
|
xp: Optional[Union[object, ArrayLibraryName]] = None,
|
|
277
339
|
dtype: Optional[DType] = None,
|
|
278
|
-
device: Optional[
|
|
340
|
+
device: Optional[AnyDevice] = None
|
|
279
341
|
):
|
|
280
342
|
"""
|
|
281
343
|
Initialize the context manager with the specified `array namespace`, `dtype` and `device`.
|
|
@@ -286,11 +348,11 @@ class array_context:
|
|
|
286
348
|
The target `array namespace` or array library name for the context.
|
|
287
349
|
- `None`: Use the `compatibility namespace` from the context;
|
|
288
350
|
|
|
289
|
-
dtype : Optional[
|
|
351
|
+
dtype : Optional[DType], default to `None`
|
|
290
352
|
The target `dtype` for the context.
|
|
291
353
|
- `None`: Use the `dtype` from the context.
|
|
292
354
|
|
|
293
|
-
device : Optional[
|
|
355
|
+
device : Optional[AnyDevice], default to `None`
|
|
294
356
|
The target `device` for the context.
|
|
295
357
|
- `None`: Use the `device` from the context.
|
|
296
358
|
"""
|
|
@@ -313,7 +375,7 @@ class array_context:
|
|
|
313
375
|
|
|
314
376
|
def unify_args(
|
|
315
377
|
ref: Optional[Union[str, int]] = 0,
|
|
316
|
-
/,
|
|
378
|
+
/, *,
|
|
317
379
|
filter_arraylike: bool = True,
|
|
318
380
|
api_version: Optional[str] = None,
|
|
319
381
|
use_compat: Optional[bool] = None,
|
|
@@ -357,6 +419,18 @@ def unify_args(
|
|
|
357
419
|
Raises
|
|
358
420
|
------
|
|
359
421
|
Refer to :func:`default.default_spec`, :func:`as_context` for possible exceptions.
|
|
422
|
+
|
|
423
|
+
Examples
|
|
424
|
+
--------
|
|
425
|
+
Use `unify_args` as a decorator to normalize array arguments before the function body runs:
|
|
426
|
+
|
|
427
|
+
>>> import numpy as np
|
|
428
|
+
>>> from cobra_array import unify_args
|
|
429
|
+
>>> @unify_args(ref=0, unify_dtype=False, unify_device=False)
|
|
430
|
+
... def add(x, y):
|
|
431
|
+
... return x + y
|
|
432
|
+
>>> add(np.asarray([1, 2]), y=np.asarray([3, 4]))
|
|
433
|
+
NumPy_Array([4 6])
|
|
360
434
|
"""
|
|
361
435
|
def decorator(func):
|
|
362
436
|
@wraps(func)
|
|
@@ -43,7 +43,7 @@ def array_namespace_alias(xp: object) -> str:
|
|
|
43
43
|
Raises
|
|
44
44
|
------
|
|
45
45
|
UnsupportedNameSpaceError
|
|
46
|
-
If the input object is not a supported namespace
|
|
46
|
+
If the input object is not a supported `array namespace`.
|
|
47
47
|
"""
|
|
48
48
|
if isinstance(xp, ModuleType):
|
|
49
49
|
if api.is_numpy_namespace(xp):
|
|
@@ -15,6 +15,24 @@ Attributes
|
|
|
15
15
|
Functions
|
|
16
16
|
---------
|
|
17
17
|
- :func:`resolve_device`: Get the device string from an object or a device specification string, and check if it is compatible with the specified `array namespace` if provided.
|
|
18
|
+
|
|
19
|
+
Examples
|
|
20
|
+
--------
|
|
21
|
+
- Basic usage::
|
|
22
|
+
|
|
23
|
+
from cobra_array.array_api import resolve_device, torch_xp, numpy_xp
|
|
24
|
+
|
|
25
|
+
r = resolve_device("cpu") # "cpu"
|
|
26
|
+
r = resolve_device("cuda:0") # "cuda:0"
|
|
27
|
+
|
|
28
|
+
if numpy_xp is not None:
|
|
29
|
+
r = resolve_device("cpu", xp="numpy")
|
|
30
|
+
|
|
31
|
+
if torch_xp is not None:
|
|
32
|
+
r = resolve_device("cpu", xp="torch")
|
|
33
|
+
|
|
34
|
+
if torch_xp is not None:
|
|
35
|
+
r = resolve_device("cpu", xp=torch_xp)
|
|
18
36
|
"""
|
|
19
37
|
|
|
20
38
|
from __future__ import annotations
|
|
@@ -61,7 +79,7 @@ def resolve_device(
|
|
|
61
79
|
obj : object
|
|
62
80
|
The input object or device specification string to extract the device information from.
|
|
63
81
|
|
|
64
|
-
xp : Optional[Union[
|
|
82
|
+
xp : Optional[Union[Any, ArrayLibraryName]], default is `None`
|
|
65
83
|
The `array namespace` to check the device compatibility against.
|
|
66
84
|
- `None`: No compatibility check will be performed.
|
|
67
85
|
|
|
@@ -80,6 +98,54 @@ def resolve_device(
|
|
|
80
98
|
If the extracted device is not compatible with the specified `array namespace`.
|
|
81
99
|
CUDAUnavailableError
|
|
82
100
|
If a CUDA device is specified but CUDA is not available for `PyTorch`.
|
|
101
|
+
|
|
102
|
+
Examples
|
|
103
|
+
--------
|
|
104
|
+
Basic parsing and normalization:
|
|
105
|
+
|
|
106
|
+
>>> resolve_device("cpu")
|
|
107
|
+
'cpu'
|
|
108
|
+
>>> resolve_device(" CUDA:0 ")
|
|
109
|
+
'cuda:0'
|
|
110
|
+
>>> resolve_device(None)
|
|
111
|
+
None
|
|
112
|
+
|
|
113
|
+
Namespace compatibility checks:
|
|
114
|
+
|
|
115
|
+
>>> resolve_device("cpu")
|
|
116
|
+
'cpu'
|
|
117
|
+
>>> resolve_device("cpu", xp="numpy")
|
|
118
|
+
'cpu'
|
|
119
|
+
>>> resolve_device("cpu", xp="torch")
|
|
120
|
+
'cpu'
|
|
121
|
+
|
|
122
|
+
Unsupported device for NumPy:
|
|
123
|
+
|
|
124
|
+
>>> resolve_device("cuda:0", xp="numpy")
|
|
125
|
+
Traceback (most recent call last):
|
|
126
|
+
...
|
|
127
|
+
cobra_array.exceptions.DeviceNotSupportedError: ...
|
|
128
|
+
|
|
129
|
+
Unsupported device type for PyTorch:
|
|
130
|
+
|
|
131
|
+
>>> resolve_device("quantum", xp="torch")
|
|
132
|
+
Traceback (most recent call last):
|
|
133
|
+
...
|
|
134
|
+
cobra_array.exceptions.DeviceNotSupportedError: ...
|
|
135
|
+
|
|
136
|
+
CUDA path for PyTorch (works with or without CUDA runtime):
|
|
137
|
+
|
|
138
|
+
>>> from cobra_array.array_api import CUDA_AVAILABLE
|
|
139
|
+
>>> result = None
|
|
140
|
+
>>> if not CUDA_AVAILABLE:
|
|
141
|
+
... try:
|
|
142
|
+
... resolve_device("cuda:0", xp="torch")
|
|
143
|
+
... except CUDAUnavailableError:
|
|
144
|
+
... result = "CUDAUnavailableError"
|
|
145
|
+
... else:
|
|
146
|
+
... result = resolve_device("cuda:0", xp="torch")
|
|
147
|
+
>>> result in {"cuda:0", "CUDAUnavailableError"}
|
|
148
|
+
True
|
|
83
149
|
"""
|
|
84
150
|
# source
|
|
85
151
|
if obj is None:
|
|
@@ -10,6 +10,23 @@ Classes
|
|
|
10
10
|
-------
|
|
11
11
|
- :class:`CompatArray`: A backend-agnostic array abstraction compliant with the `Python Array API standard`.
|
|
12
12
|
- :class:`CompatNamespace`: A wrapper around an `array namespace` providing a unified, backend-agnostic functional interface.
|
|
13
|
+
|
|
14
|
+
Examples
|
|
15
|
+
--------
|
|
16
|
+
- Basic usage::
|
|
17
|
+
|
|
18
|
+
import numpy as np
|
|
19
|
+
from cobra_array.compat import CompatArray, CompatNamespace, wrap_arraylike, unwrap
|
|
20
|
+
|
|
21
|
+
cxp = CompatNamespace(np)
|
|
22
|
+
a = CompatArray(np.asarray([1, 2, 3]))
|
|
23
|
+
b = cxp.asarray([10, 20, 30])
|
|
24
|
+
|
|
25
|
+
r = (a + b).to_list() # [11, 22, 33]
|
|
26
|
+
r = cxp.add(a, b).to_list() # [11, 22, 33]
|
|
27
|
+
|
|
28
|
+
wrapped = wrap_arraylike(np.asarray([4, 5]))
|
|
29
|
+
unwrap(wrapped).tolist() # [4, 5]
|
|
13
30
|
"""
|
|
14
31
|
|
|
15
32
|
from ._array import (CompatArray, wrap_arraylike, unwrap)
|
|
@@ -26,6 +26,61 @@ class CompatArray(Compat):
|
|
|
26
26
|
- All operations follow the semantics defined by the `Python Array API standard`.
|
|
27
27
|
- Methods correspond directly to standard functions, but are exposed in an object-oriented form.
|
|
28
28
|
- All methods guarantee that any array-like objects in the returned value are automatically wrapped as :class:`CompatArray`. This applies recursively to arrays contained in Python containers (e.g., `tuple`, `list`, `dict`). Non-array objects remain unchanged.
|
|
29
|
+
|
|
30
|
+
Examples
|
|
31
|
+
--------
|
|
32
|
+
Create a :class:`CompatArray` from a backend array object:
|
|
33
|
+
|
|
34
|
+
>>> xp = to_xp("numpy")
|
|
35
|
+
>>> a = CompatArray(xp.asarray([1, 2, 3]))
|
|
36
|
+
>>> a
|
|
37
|
+
NumPy_Array([1 2 3])
|
|
38
|
+
>>> a.to_list()
|
|
39
|
+
[1, 2, 3]
|
|
40
|
+
>>> a.shape
|
|
41
|
+
(3,)
|
|
42
|
+
>>> a.xp_name
|
|
43
|
+
'NumPy'
|
|
44
|
+
|
|
45
|
+
Convert Python data with an explicit backend:
|
|
46
|
+
|
|
47
|
+
>>> b = CompatArray.from_other([10, 20], xp="numpy")
|
|
48
|
+
>>> isinstance(b, CompatArray)
|
|
49
|
+
True
|
|
50
|
+
>>> b.to_tensor()
|
|
51
|
+
tensor([10, 20])
|
|
52
|
+
|
|
53
|
+
Re-wrapping behavior for existing :class:`CompatArray`:
|
|
54
|
+
|
|
55
|
+
>>> CompatArray(b) is b
|
|
56
|
+
True
|
|
57
|
+
>>> CompatArray(b, copy=True) is b
|
|
58
|
+
False
|
|
59
|
+
|
|
60
|
+
Call elementwise functions directly:
|
|
61
|
+
|
|
62
|
+
>>> x = CompatArray.from_other([1, 2, 3], xp="numpy")
|
|
63
|
+
>>> y = CompatArray.from_other([10, 20, 30], xp="numpy")
|
|
64
|
+
>>> x.add(y)
|
|
65
|
+
NumPy_Array([11 22 33])
|
|
66
|
+
>>> x.multiply(2)
|
|
67
|
+
NumPy_Array([2 4 6])
|
|
68
|
+
|
|
69
|
+
Use Python operators:
|
|
70
|
+
|
|
71
|
+
>>> x + y
|
|
72
|
+
NumPy_Array([11 22 33])
|
|
73
|
+
>>> y - x
|
|
74
|
+
NumPy_Array([ 9 18 27])
|
|
75
|
+
>>> x * 2
|
|
76
|
+
NumPy_Array([2 4 6])
|
|
77
|
+
|
|
78
|
+
Invalid input raises an exception:
|
|
79
|
+
|
|
80
|
+
>>> CompatArray("not-array")
|
|
81
|
+
Traceback (most recent call last):
|
|
82
|
+
...
|
|
83
|
+
cobra_array.exceptions.NotArrayAPIObjectError: ...
|
|
29
84
|
"""
|
|
30
85
|
_arr = None
|
|
31
86
|
_cxp = None
|
|
@@ -103,7 +158,7 @@ class CompatArray(Compat):
|
|
|
103
158
|
|
|
104
159
|
Parameters
|
|
105
160
|
----------
|
|
106
|
-
device :
|
|
161
|
+
device : AnyDevice
|
|
107
162
|
A device object or name.
|
|
108
163
|
|
|
109
164
|
stream : Optional[Union[int, Any]], default to `None`
|
|
@@ -9,14 +9,14 @@ from typing import (Union, List, Tuple, Optional, Any, Sequence, Generic, TypeVa
|
|
|
9
9
|
|
|
10
10
|
from ._base import Compat
|
|
11
11
|
from ..types import (
|
|
12
|
-
T, DTypeT, DeviceT, dtypeT, deviceT, DType,
|
|
12
|
+
T, DTypeT, DeviceT, dtypeT, deviceT, DType, AnyDevice,
|
|
13
13
|
ArrayLike, ArrayLibraryName,
|
|
14
14
|
ArrayOrAny, ArrayOrScalar, ArrayOrReal, ArrayOrIntLike, ArrayOrInt, ArrayOrbool,
|
|
15
15
|
UniqueAllResult, UniqueCountsResult, UniqueInverseResult
|
|
16
16
|
)
|
|
17
17
|
|
|
18
18
|
TT = TypeVar("TT", bound=DType)
|
|
19
|
-
DT = TypeVar("DT", bound=
|
|
19
|
+
DT = TypeVar("DT", bound=AnyDevice)
|
|
20
20
|
|
|
21
21
|
|
|
22
22
|
class CompatArray(Compat, Generic[TT, DT]):
|
|
@@ -31,47 +31,54 @@ class CompatArray(Compat, Generic[TT, DT]):
|
|
|
31
31
|
def from_other(cls, obj: object, /, xp: Literal["numpy"]) -> CompatArray[Any, Literal["cpu"]]: ...
|
|
32
32
|
@overload
|
|
33
33
|
@classmethod
|
|
34
|
-
def from_other(cls, obj: NDArray[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT,
|
|
34
|
+
def from_other(cls, obj: NDArray[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, AnyDevice]: ...
|
|
35
35
|
@overload
|
|
36
36
|
@classmethod
|
|
37
|
-
def from_other(cls, obj: ArrayLike[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT,
|
|
37
|
+
def from_other(cls, obj: ArrayLike[dtypeT], /, xp: Union[object, ArrayLibraryName]) -> CompatArray[dtypeT, AnyDevice]: ...
|
|
38
38
|
@overload
|
|
39
39
|
@classmethod
|
|
40
|
-
def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any,
|
|
40
|
+
def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, AnyDevice]: ...
|
|
41
41
|
@classmethod
|
|
42
|
-
def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any,
|
|
42
|
+
def from_other(cls, obj: object, /, xp: Union[object, ArrayLibraryName]) -> CompatArray[Any, AnyDevice]: ...
|
|
43
43
|
|
|
44
44
|
@overload
|
|
45
45
|
def __new__(cls, arr: NDArray[dtypeT], /, **kwargs) -> CompatArray[dtypeT, Literal["cpu"]]: ...
|
|
46
46
|
@overload
|
|
47
|
-
def __new__(cls, arr: Tensor, /, **kwargs) -> CompatArray[Any,
|
|
47
|
+
def __new__(cls, arr: Tensor, /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
|
|
48
48
|
@overload
|
|
49
49
|
def __new__(cls, arr: CompatArray[dtypeT, deviceT], /, **kwargs) -> CompatArray[dtypeT, deviceT]: ...
|
|
50
50
|
@overload
|
|
51
|
-
def __new__(cls, arr: ArrayLike[dtypeT], /, **kwargs) -> CompatArray[dtypeT,
|
|
52
|
-
def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any,
|
|
51
|
+
def __new__(cls, arr: ArrayLike[dtypeT], /, **kwargs) -> CompatArray[dtypeT, AnyDevice]: ...
|
|
52
|
+
def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any, AnyDevice]: ...
|
|
53
53
|
|
|
54
54
|
# === Conversion functions ===
|
|
55
55
|
def to_numpy(self, *, copy: bool = False) -> NDArray[TT]: ...
|
|
56
|
-
def to_tensor(self, *, device: Optional[
|
|
56
|
+
def to_tensor(self, *, device: Optional[AnyDevice] = None, copy: bool = False) -> Tensor: ...
|
|
57
57
|
def to_list(self, *, copy: bool = False) -> List[TT]: ...
|
|
58
58
|
|
|
59
59
|
# === Device functions ===
|
|
60
|
+
@overload
|
|
60
61
|
def to_device(self, device: DeviceT, /, *, stream: Optional[Any] = None) -> CompatArray[TT, DeviceT]: ...
|
|
62
|
+
@overload
|
|
63
|
+
def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[TT, AnyDevice]: ...
|
|
64
|
+
|
|
65
|
+
def to_device(self, device: AnyDevice, /, *, stream: Optional[Any] = None) -> CompatArray[Any, AnyDevice]: ...
|
|
61
66
|
|
|
62
67
|
# === Data type functions ===
|
|
63
68
|
@overload
|
|
64
|
-
def astype(self, dtype: DTypeT, /, *, copy: bool = ...) -> CompatArray[DTypeT, DT]: ...
|
|
69
|
+
def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: None = ...) -> CompatArray[DTypeT, DT]: ...
|
|
65
70
|
@overload
|
|
66
71
|
def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: DeviceT) -> CompatArray[DTypeT, DeviceT]: ...
|
|
72
|
+
@overload
|
|
73
|
+
def astype(self, dtype: DTypeT, /, *, copy: bool = ..., device: AnyDevice) -> CompatArray[DTypeT, AnyDevice]: ...
|
|
67
74
|
|
|
68
75
|
def astype(
|
|
69
76
|
self,
|
|
70
77
|
dtype: DType,
|
|
71
78
|
/, *,
|
|
72
79
|
copy: bool = True,
|
|
73
|
-
device: Optional[
|
|
74
|
-
) -> CompatArray[Any,
|
|
80
|
+
device: Optional[AnyDevice] = None
|
|
81
|
+
) -> CompatArray[Any, AnyDevice]:
|
|
75
82
|
"""
|
|
76
83
|
Copies `self` to a specified data type irrespective of Type Promotion Rules rules.
|
|
77
84
|
|
|
@@ -88,8 +95,8 @@ class CompatArray(Compat, Generic[TT, DT]):
|
|
|
88
95
|
- `True`: `self` must be returned;
|
|
89
96
|
- `False`: A newly allocated array must be returned.
|
|
90
97
|
|
|
91
|
-
device : Optional[
|
|
92
|
-
|
|
98
|
+
device : Optional[AnyDevice], default to `None`
|
|
99
|
+
Device on which to place the returned array.
|
|
93
100
|
- `None`: The output array device must be inferred from `self`.
|
|
94
101
|
|
|
95
102
|
Returns
|
|
@@ -1793,7 +1800,7 @@ class CompatArray(Compat, Generic[TT, DT]):
|
|
|
1793
1800
|
|
|
1794
1801
|
|
|
1795
1802
|
@overload
|
|
1796
|
-
def unwrap(obj: CompatArray[dtypeT,
|
|
1803
|
+
def unwrap(obj: CompatArray[dtypeT, AnyDevice]) -> ArrayLike[dtypeT]: ...
|
|
1797
1804
|
@overload
|
|
1798
1805
|
def unwrap(obj: ArrayLike[dtypeT]) -> ArrayLike[dtypeT]: ...
|
|
1799
1806
|
@overload
|
|
@@ -1808,7 +1815,7 @@ def wrap_arraylike(arr: NDArray[dtypeT], xp: Optional[object] = ...) -> CompatAr
|
|
|
1808
1815
|
@overload
|
|
1809
1816
|
def wrap_arraylike(arr: CompatArray[dtypeT, deviceT], xp: Optional[object] = ...) -> CompatArray[dtypeT, deviceT]: ...
|
|
1810
1817
|
@overload
|
|
1811
|
-
def wrap_arraylike(arr: ArrayLike[dtypeT], xp: Optional[object] = ...) -> CompatArray[dtypeT,
|
|
1818
|
+
def wrap_arraylike(arr: ArrayLike[dtypeT], xp: Optional[object] = ...) -> CompatArray[dtypeT, AnyDevice]: ...
|
|
1812
1819
|
@overload
|
|
1813
1820
|
def wrap_arraylike(arr: T, xp: Optional[object] = ...) -> T: ...
|
|
1814
1821
|
|
|
@@ -23,6 +23,31 @@ class CompatNamespace(Compat):
|
|
|
23
23
|
- Functions correspond directly to those defined in the underlying `array namespace`, following the `Python Array API standard`.
|
|
24
24
|
- This namespace complements :class:`CompatArray` by providing a functional interface for operations that are not naturally expressed as methods.
|
|
25
25
|
- All functions guarantee that any array-like objects in the returned value are automatically wrapped as :class:`CompatArray`. This conversion is applied recursively to arrays contained in Python containers (e.g., `tuple`, `list`, `dict`). Non-array objects remain unchanged.
|
|
26
|
+
|
|
27
|
+
Examples
|
|
28
|
+
--------
|
|
29
|
+
Create a compatibility namespace from a backend namespace:
|
|
30
|
+
|
|
31
|
+
>>> import numpy as np
|
|
32
|
+
>>> cxp = CompatNamespace(np)
|
|
33
|
+
>>> cxp.xp_name
|
|
34
|
+
'NumPy'
|
|
35
|
+
|
|
36
|
+
Create arrays and call namespace functions:
|
|
37
|
+
|
|
38
|
+
>>> a = cxp.asarray([1, 2, 3])
|
|
39
|
+
>>> b = cxp.asarray([10, 20, 30])
|
|
40
|
+
>>> a
|
|
41
|
+
NumPy_Array([1 2 3])
|
|
42
|
+
>>> cxp.add(a, b).to_list()
|
|
43
|
+
[11, 22, 33]
|
|
44
|
+
|
|
45
|
+
Missing attributes raise an exception:
|
|
46
|
+
|
|
47
|
+
>>> cxp.this_attr_does_not_exist
|
|
48
|
+
Traceback (most recent call last):
|
|
49
|
+
...
|
|
50
|
+
AttributeError: ...
|
|
26
51
|
"""
|
|
27
52
|
def __new__(cls, xp, /):
|
|
28
53
|
if isinstance(xp, CompatNamespace):
|