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.
Files changed (31) hide show
  1. {cobra_array-0.1.0/src/cobra_array.egg-info → cobra_array-0.1.2}/PKG-INFO +3 -11
  2. {cobra_array-0.1.0 → cobra_array-0.1.2}/README.md +2 -10
  3. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/__init__.py +7 -15
  4. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/_core.py +84 -10
  5. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/_utils.py +1 -1
  6. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/array_api.py +67 -1
  7. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/__init__.py +17 -0
  8. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_array.py +56 -1
  9. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_array.pyi +24 -17
  10. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_namespace.py +25 -0
  11. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_namespace.pyi +122 -70
  12. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/convert.py +53 -0
  13. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/convert.pyi +13 -13
  14. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/default.py +30 -11
  15. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/types.py +17 -15
  16. {cobra_array-0.1.0 → cobra_array-0.1.2/src/cobra_array.egg-info}/PKG-INFO +3 -11
  17. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/SOURCES.txt +2 -0
  18. cobra_array-0.1.2/tests/test_compat.py +253 -0
  19. cobra_array-0.1.2/tests/test_compat_namespace.py +122 -0
  20. {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_default.py +1 -1
  21. {cobra_array-0.1.0 → cobra_array-0.1.2}/LICENSE +0 -0
  22. {cobra_array-0.1.0 → cobra_array-0.1.2}/pyproject.toml +0 -0
  23. {cobra_array-0.1.0 → cobra_array-0.1.2}/setup.cfg +0 -0
  24. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/compat/_base.py +0 -0
  25. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array/exceptions.py +0 -0
  26. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/dependency_links.txt +0 -0
  27. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/requires.txt +0 -0
  28. {cobra_array-0.1.0 → cobra_array-0.1.2}/src/cobra_array.egg-info/top_level.txt +0 -0
  29. {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_backend.py +0 -0
  30. {cobra_array-0.1.0 → cobra_array-0.1.2}/tests/test_convert.py +0 -0
  31. {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.0
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.0"
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, Device, ArrayLike, ArrayLibraryName)
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, Any]: ...
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, 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, 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, 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[Device] = None
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[DTypeT], default to `None`
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[DeviceT], default to `None`
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[Namespace, CompatNamespace, ArrayLibraryName]], default is `None`
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 : 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, Device,
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=Device)
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, Any]: ...
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, Any]: ...
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, 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, 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, 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, Any]: ...
52
- def __new__(cls, arr: ArrayLike[Any], /, **kwargs) -> CompatArray[Any, 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[DeviceT] = None, copy: bool = False) -> Tensor: ...
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[Device] = None
74
- ) -> CompatArray[Any, 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[DeviceT], default to `None`
92
- DeviceT on which to place the returned array.
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, Any]) -> ArrayLike[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, Any]: ...
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):