cunumpy 0.1.2__tar.gz → 0.1.4__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cunumpy
3
- Version: 0.1.2
3
+ Version: 0.1.4
4
4
  Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
5
5
  Author: Max
6
6
  Project-URL: Source, https://github.com/max-models/cunumpy
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
16
  Requires-Python: >=3.8
17
17
  Description-Content-Type: text/markdown
18
+ Requires-Dist: array-api-compat
18
19
  Requires-Dist: numpy
19
20
  Provides-Extra: dev
20
21
  Requires-Dist: black[jupyter]; extra == "dev"
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
32
33
  Requires-Dist: sphinx-book-theme; extra == "docs"
33
34
  Provides-Extra: test
34
35
  Requires-Dist: coverage; extra == "test"
36
+ Requires-Dist: pyccel; extra == "test"
35
37
  Requires-Dist: pytest; extra == "test"
36
38
 
37
39
  # CuNumpy
@@ -52,7 +54,8 @@ export ARRAY_BACKEND=cupy
52
54
 
53
55
  ```python
54
56
  import cunumpy as xp
55
- arr = xp.array([1,2])
57
+
58
+ arr = xp.array([1, 2])
56
59
 
57
60
  print(type(arr))
58
61
  print(xp.__version__)
@@ -16,7 +16,8 @@ export ARRAY_BACKEND=cupy
16
16
 
17
17
  ```python
18
18
  import cunumpy as xp
19
- arr = xp.array([1,2])
19
+
20
+ arr = xp.array([1, 2])
20
21
 
21
22
  print(type(arr))
22
23
  print(xp.__version__)
@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
5
5
 
6
6
  [project]
7
7
  name = "cunumpy"
8
- version = "0.1.2"
8
+ version = "0.1.4"
9
9
  description = "Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`."
10
10
  readme = "README.md"
11
11
  keywords = [ "python" ]
@@ -23,6 +23,7 @@ classifiers = [
23
23
  "Programming Language :: Python :: 3.13",
24
24
  ]
25
25
  dependencies = [
26
+ "array-api-compat",
26
27
  "numpy",
27
28
  ]
28
29
 
@@ -43,7 +44,7 @@ optional-dependencies.docs = [
43
44
  "sphinx",
44
45
  "sphinx-book-theme",
45
46
  ]
46
- optional-dependencies.test = [ "coverage", "pytest" ]
47
+ optional-dependencies.test = [ "coverage", "pyccel", "pytest" ]
47
48
  urls."Source" = "https://github.com/max-models/cunumpy"
48
49
 
49
50
  [tool.setuptools.packages.find]
@@ -1,10 +1,15 @@
1
1
  # cunumpy/__init__.py
2
+ from importlib.metadata import PackageNotFoundError, version
3
+
2
4
  from . import xp
5
+ from .kernel import PyccelKernel
3
6
  from .xp import (
7
+ cupy_available,
4
8
  get_backend,
5
9
  is_cpu,
6
10
  is_gpu,
7
11
  set_backend,
12
+ set_device,
8
13
  synchronize,
9
14
  to_cunumpy,
10
15
  to_cupy,
@@ -12,19 +17,28 @@ from .xp import (
12
17
  use_backend,
13
18
  )
14
19
 
20
+ try:
21
+ __version__ = version("cunumpy")
22
+ except PackageNotFoundError:
23
+ __version__ = "0.0.0+unknown"
24
+
15
25
  __all__ = [
16
- "xp",
17
- "to_numpy",
18
- "to_cupy",
19
- "to_cunumpy",
26
+ "PyccelKernel",
27
+ "__version__",
28
+ "cupy_available",
29
+ "cupy_backend",
20
30
  "get_backend",
21
- "is_gpu",
22
31
  "is_cpu",
23
- "use_backend",
32
+ "is_gpu",
33
+ "numpy_backend",
24
34
  "set_backend",
35
+ "set_device",
25
36
  "synchronize",
26
- "numpy_backend",
27
- "cupy_backend",
37
+ "to_cunumpy",
38
+ "to_cupy",
39
+ "to_numpy",
40
+ "use_backend",
41
+ "xp",
28
42
  ]
29
43
 
30
44
 
@@ -7,18 +7,22 @@ from typing import Any, Generator
7
7
  import numpy as np
8
8
  from numpy import *
9
9
 
10
- from . import xp
10
+ from . import xp as xp
11
+ from .kernel import PyccelKernel as PyccelKernel
11
12
 
12
13
  def to_numpy(array: Any) -> np.ndarray: ...
13
14
  def to_cupy(array: Any) -> Any: ...
14
15
  def to_cunumpy(array: Any) -> Any: ...
16
+ def cupy_available() -> bool: ...
15
17
  def get_backend(array: Any) -> str: ...
16
18
  def is_gpu(array: Any) -> bool: ...
17
19
  def is_cpu(array: Any) -> bool: ...
18
20
  @contextmanager
19
- def use_backend(backend: str) -> Generator[None, None, None]: ...
21
+ def use_backend(backend: str) -> Generator[None]: ...
20
22
  def set_backend(backend: str) -> None: ...
23
+ def set_device(device_id: int) -> None: ...
21
24
  def synchronize() -> None: ...
22
25
 
23
26
  numpy_backend: bool
24
27
  cupy_backend: bool
28
+ __version__: str
@@ -0,0 +1,344 @@
1
+ """Interface for calling Pyccel-compiled kernels with CuPy arrays.
2
+
3
+ Kernels generated by `pyccel <https://github.com/pyccel/pyccel>`_ are compiled
4
+ C/Fortran routines that only understand NumPy (host) arrays. :class:`PyccelKernel`
5
+ wraps such a kernel so that it can be called transparently with CuPy (device)
6
+ arrays: the arguments are copied to the host before the call, in-place updates
7
+ made by the kernel are copied back to the device afterwards, and any arrays
8
+ returned by the kernel are moved back to the device.
9
+
10
+ On the NumPy backend the wrapper is a no-op and the kernel is called directly.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import copy
16
+ from typing import Any, Callable, Sequence
17
+
18
+ import array_api_compat
19
+ import numpy as np
20
+
21
+ from .xp import _cupy_backend, to_cupy, to_numpy
22
+
23
+ __all__ = ["PyccelKernel"]
24
+
25
+
26
+ class PyccelKernel:
27
+ """Call a Pyccel-compiled kernel with NumPy or CuPy arrays.
28
+
29
+ Parameters
30
+ ----------
31
+ kernel : callable
32
+ The pyccelized kernel (or any callable expecting NumPy arrays).
33
+ use_cupy : bool, optional
34
+ Force host/device conversion on (``True``) or off (``False``). By
35
+ default (``None``) it is decided at call time: conversion happens when
36
+ the active backend is CuPy or when a CuPy array is passed in.
37
+ object_modules : sequence of str, optional
38
+ Module prefixes (e.g. ``("struphy.", "feectools.")``) whose instances
39
+ should be traversed attribute-by-attribute when looking for arrays to
40
+ convert. Objects from other modules are passed through untouched.
41
+ outputs : sequence of int or str, optional
42
+ Which arguments the kernel writes to. Only those are copied back to the
43
+ device after the call, which avoids pointless device transfers for the
44
+ (usually much larger) read-only inputs. Positional arguments are named
45
+ by index, keyword arguments by name::
46
+
47
+ interpolate = PyccelKernel(some_interpolation_kernel, outputs=(5,))
48
+ interpolate(x, y, z, basis, coeffs, out) # `out` is argument 5
49
+
50
+ Pyccel-compiled kernels are builtins with no introspectable signature,
51
+ so an index and a name are *not* interchangeable: declare the form you
52
+ actually call with. An empty sequence declares that the kernel writes to
53
+ none of its arguments. By default (``None``) every converted array is
54
+ copied back, which is always correct but does more work.
55
+
56
+ Examples
57
+ --------
58
+ >>> from cunumpy.kernel import PyccelKernel
59
+ >>> kernel = PyccelKernel(my_pyccelized_function)
60
+ >>> kernel(out, x, y) # `out`, `x`, `y` may be NumPy or CuPy arrays
61
+ """
62
+
63
+ def __init__(
64
+ self,
65
+ kernel: Callable[..., Any],
66
+ use_cupy: bool | None = None,
67
+ object_modules: Sequence[str] = (),
68
+ outputs: Sequence[int | str] | None = None,
69
+ ) -> None:
70
+ self._kernel = kernel
71
+ self._use_cupy = use_cupy
72
+ self._object_modules = tuple(object_modules)
73
+
74
+ if outputs is None:
75
+ self._outputs: tuple[int | str, ...] | None = None
76
+ else:
77
+ if isinstance(outputs, (int, str)):
78
+ raise TypeError(
79
+ "outputs must be a sequence of argument indices/names, "
80
+ f"not a bare {type(outputs).__name__} "
81
+ f"(did you mean outputs=({outputs!r},)?)"
82
+ )
83
+ for entry in outputs:
84
+ if not isinstance(entry, (int, str)) or isinstance(entry, bool):
85
+ raise TypeError(
86
+ "outputs entries must be argument indices (int) or "
87
+ f"names (str), got {entry!r}"
88
+ )
89
+ self._outputs = tuple(outputs)
90
+
91
+ def __repr__(self) -> str:
92
+ return (
93
+ f"PyccelKernel(kernel={self.name!r}, use_cupy={self.use_cupy!r}, "
94
+ f"outputs={self._outputs!r})"
95
+ )
96
+
97
+ def _convert_to_numpy(
98
+ self,
99
+ value: Any,
100
+ converted: list[tuple[Any, np.ndarray]],
101
+ memo: dict[int, Any],
102
+ ) -> Any:
103
+ """Recursively replace CuPy arrays in `value` by host copies.
104
+
105
+ Every replacement is appended to `converted` as a
106
+ ``(device_array, host_copy)`` pair.
107
+
108
+ `memo` maps ``id(original) -> converted`` and is shared across all
109
+ arguments of a single call. It serves two purposes: a device array
110
+ reachable by several paths is copied to the host exactly once (so the
111
+ kernel sees one shared array, as the caller intended, and the write-back
112
+ happens once), and reference cycles terminate instead of recursing
113
+ forever. Everything traversed here stays reachable from the caller's
114
+ arguments for the duration of the call, so the `id` keys cannot be
115
+ reused by unrelated objects.
116
+ """
117
+ key = id(value)
118
+ if key in memo:
119
+ return memo[key]
120
+
121
+ if array_api_compat.is_cupy_array(value):
122
+ value_np = to_numpy(value)
123
+ memo[key] = value_np
124
+ converted.append((value, value_np))
125
+ return value_np
126
+
127
+ if isinstance(value, tuple):
128
+ # A tuple cannot be memoized before its items are converted, but it
129
+ # can only take part in a cycle through a mutable container, and
130
+ # those are memoized before they are filled in below.
131
+ value_np = tuple(
132
+ self._convert_to_numpy(item, converted, memo) for item in value
133
+ )
134
+ memo[key] = value_np
135
+ return value_np
136
+
137
+ if isinstance(value, list):
138
+ value_np = []
139
+ memo[key] = value_np
140
+ value_np.extend(
141
+ self._convert_to_numpy(item, converted, memo) for item in value
142
+ )
143
+ return value_np
144
+
145
+ if isinstance(value, dict):
146
+ value_np = {}
147
+ memo[key] = value_np
148
+ for k, v in value.items():
149
+ value_np[k] = self._convert_to_numpy(v, converted, memo)
150
+ return value_np
151
+
152
+ if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
153
+ self._object_modules
154
+ ):
155
+ # Shallow-copy the object so the caller's instance keeps pointing at
156
+ # its device arrays; only the copy holds the host views.
157
+ value_np = copy.copy(value)
158
+ memo[key] = value_np
159
+ for name, attr in vars(value).items():
160
+ setattr(value_np, name, self._convert_to_numpy(attr, converted, memo))
161
+ return value_np
162
+
163
+ return value
164
+
165
+ @staticmethod
166
+ def _convert_from_numpy(value: Any) -> Any:
167
+ """Move NumPy arrays returned by the kernel back to the device."""
168
+ if isinstance(value, np.ndarray):
169
+ return to_cupy(value)
170
+ if isinstance(value, tuple):
171
+ return tuple(PyccelKernel._convert_from_numpy(item) for item in value)
172
+ if isinstance(value, list):
173
+ return [PyccelKernel._convert_from_numpy(item) for item in value]
174
+ return value
175
+
176
+ def _collect_host_arrays(self, value: Any, found: set[int], seen: set[int]) -> None:
177
+ """Record the id of every host array reachable from `value`.
178
+
179
+ Runs over the *converted* arguments, using the same traversal rules as
180
+ :meth:`_convert_to_numpy`, so that an output declared as a container or
181
+ an object contributes the arrays nested inside it.
182
+ """
183
+ if isinstance(value, np.ndarray):
184
+ found.add(id(value))
185
+ return
186
+
187
+ if id(value) in seen:
188
+ return
189
+
190
+ if isinstance(value, (tuple, list)):
191
+ seen.add(id(value))
192
+ for item in value:
193
+ self._collect_host_arrays(item, found, seen)
194
+ return
195
+
196
+ if isinstance(value, dict):
197
+ seen.add(id(value))
198
+ for item in value.values():
199
+ self._collect_host_arrays(item, found, seen)
200
+ return
201
+
202
+ if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
203
+ self._object_modules
204
+ ):
205
+ seen.add(id(value))
206
+ for attr in vars(value).values():
207
+ self._collect_host_arrays(attr, found, seen)
208
+
209
+ def _output_host_arrays(
210
+ self, args_np: list[Any], kwargs_np: dict[str, Any]
211
+ ) -> set[int]:
212
+ """Ids of the host arrays reachable from the declared output arguments.
213
+
214
+ Raises
215
+ ------
216
+ IndexError, KeyError
217
+ If a declared output does not correspond to an argument of this
218
+ call -- typically because an argument declared by index was passed
219
+ as a keyword, or vice versa.
220
+ """
221
+ found: set[int] = set()
222
+ seen: set[int] = set()
223
+
224
+ for entry in self._outputs or ():
225
+ if isinstance(entry, int):
226
+ index = entry + len(args_np) if entry < 0 else entry
227
+ if not 0 <= index < len(args_np):
228
+ raise IndexError(
229
+ f"{self.name}() was declared with output argument "
230
+ f"{entry}, but was called with {len(args_np)} "
231
+ "positional argument(s). Note that an output passed as "
232
+ "a keyword must be declared by name, not by index."
233
+ )
234
+ self._collect_host_arrays(args_np[index], found, seen)
235
+ else:
236
+ if entry not in kwargs_np:
237
+ raise KeyError(
238
+ f"{self.name}() was declared with output argument "
239
+ f"{entry!r}, but no such keyword argument was passed. "
240
+ "Note that an output passed positionally must be "
241
+ "declared by index, not by name."
242
+ )
243
+ self._collect_host_arrays(kwargs_np[entry], found, seen)
244
+
245
+ return found
246
+
247
+ def _contains_cupy(self, value: Any, seen: set[int] | None = None) -> bool:
248
+ """Whether `value` holds a CuPy array, following the same traversal
249
+ rules as :meth:`_convert_to_numpy`.
250
+
251
+ `seen` tracks already-visited containers so that reference cycles
252
+ terminate.
253
+ """
254
+ if array_api_compat.is_cupy_array(value):
255
+ return True
256
+
257
+ if seen is None:
258
+ seen = set()
259
+ if id(value) in seen:
260
+ return False
261
+
262
+ if isinstance(value, (tuple, list)):
263
+ seen.add(id(value))
264
+ return any(self._contains_cupy(item, seen) for item in value)
265
+
266
+ if isinstance(value, dict):
267
+ seen.add(id(value))
268
+ return any(self._contains_cupy(item, seen) for item in value.values())
269
+
270
+ if hasattr(value, "__dict__") and value.__class__.__module__.startswith(
271
+ self._object_modules
272
+ ):
273
+ seen.add(id(value))
274
+ return any(self._contains_cupy(attr, seen) for attr in vars(value).values())
275
+
276
+ return False
277
+
278
+ def _needs_conversion(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool:
279
+ if self._use_cupy is not None:
280
+ return self._use_cupy
281
+ if _cupy_backend():
282
+ return True
283
+ # The backend is NumPy, but individual CuPy arrays may still have been
284
+ # passed in explicitly.
285
+ return any(self._contains_cupy(value) for value in (*args, *kwargs.values()))
286
+
287
+ def __call__(self, *args: Any, **kwargs: Any) -> Any:
288
+ if not self._needs_conversion(args, kwargs):
289
+ return self._kernel(*args, **kwargs)
290
+
291
+ # Convert CuPy arrays in args/kwargs to NumPy arrays on the host. The
292
+ # memo is shared across args and kwargs so that an array passed more
293
+ # than once stays a single array on the host too.
294
+ converted: list[tuple[Any, np.ndarray]] = []
295
+ memo: dict[int, Any] = {}
296
+ args_np = [self._convert_to_numpy(x, converted, memo) for x in args]
297
+ kwargs_np = {
298
+ k: self._convert_to_numpy(v, converted, memo) for k, v in kwargs.items()
299
+ }
300
+
301
+ # Which arrays the kernel may have written to is resolved before the
302
+ # call, so a mis-declared output is reported even if the kernel itself
303
+ # would have raised first.
304
+ writeable = (
305
+ None
306
+ if self._outputs is None
307
+ else self._output_host_arrays(args_np, kwargs_np)
308
+ )
309
+
310
+ result = self._kernel(*args_np, **kwargs_np)
311
+
312
+ # Copy in-place kernel updates back to the device arrays.
313
+ for device_array, host_array in converted:
314
+ if writeable is None or id(host_array) in writeable:
315
+ device_array[...] = to_cupy(host_array)
316
+
317
+ return self._convert_from_numpy(result)
318
+
319
+ @property
320
+ def name(self) -> str:
321
+ """Name of the wrapped kernel."""
322
+ return getattr(self._kernel, "__name__", type(self._kernel).__name__)
323
+
324
+ @property
325
+ def kernel(self) -> Callable[..., Any]:
326
+ """The wrapped kernel."""
327
+ return self._kernel
328
+
329
+ @property
330
+ def use_cupy(self) -> bool:
331
+ """Whether calls currently convert between device and host arrays."""
332
+ if self._use_cupy is not None:
333
+ return self._use_cupy
334
+ return _cupy_backend()
335
+
336
+ @property
337
+ def object_modules(self) -> tuple[str, ...]:
338
+ """Module prefixes whose instances are traversed for arrays."""
339
+ return self._object_modules
340
+
341
+ @property
342
+ def outputs(self) -> tuple[int | str, ...] | None:
343
+ """Declared output arguments, or ``None`` if every array is copied back."""
344
+ return self._outputs
@@ -1,23 +1,51 @@
1
1
  import os
2
+ import warnings
2
3
  from contextlib import contextmanager
3
4
  from types import ModuleType
4
5
  from typing import TYPE_CHECKING, Any, Generator, Literal
5
6
 
6
- import numpy as np
7
+ import array_api_compat
8
+ import array_api_compat.numpy as np
7
9
 
8
10
  BackendType = Literal["numpy", "cupy"]
9
11
 
10
12
 
13
+ _CUPY_AVAILABLE_CACHE = None
14
+
15
+
16
+ def cupy_available() -> bool:
17
+ """Check if CuPy is available and functional."""
18
+ global _CUPY_AVAILABLE_CACHE
19
+ if _CUPY_AVAILABLE_CACHE is not None:
20
+ return _CUPY_AVAILABLE_CACHE
21
+
22
+ try:
23
+ import cupy as cp
24
+
25
+ # Check if a GPU is available
26
+ _CUPY_AVAILABLE_CACHE = cp.is_available()
27
+ return _CUPY_AVAILABLE_CACHE
28
+ except Exception: # noqa: BLE001 - tolerate any driver/runtime failure
29
+ _CUPY_AVAILABLE_CACHE = False
30
+ return False
31
+
32
+
11
33
  class ArrayBackend:
34
+ """Holds the process-wide active backend (NumPy or CuPy).
35
+
36
+ Not thread-safe: `set_backend`/`use_backend` mutate this single shared
37
+ instance in place, so concurrent code (threads, async tasks) switching
38
+ backends independently will race. Safe for the typical single-threaded
39
+ script/notebook usage this library targets.
40
+ """
41
+
12
42
  def __init__(
13
43
  self,
14
44
  backend: BackendType = "numpy",
15
45
  verbose: bool = False,
16
46
  ) -> None:
17
- assert backend.lower() in [
18
- "numpy",
19
- "cupy",
20
- ], "Array backend must be either 'numpy' or 'cupy'."
47
+ if backend.lower() not in ("numpy", "cupy"):
48
+ raise ValueError("Array backend must be either 'numpy' or 'cupy'.")
21
49
 
22
50
  self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy"
23
51
  self._xp: ModuleType = np # Placeholder
@@ -27,24 +55,23 @@ class ArrayBackend:
27
55
 
28
56
  def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleType:
29
57
  if backend == "cupy":
30
- try:
31
- import cupy as cp
58
+ if cupy_available():
59
+ import array_api_compat.cupy as cp
32
60
 
61
+ self._backend = "cupy"
33
62
  return cp
34
- except ImportError:
63
+ else:
35
64
  if verbose:
36
- print("CuPy not available.")
65
+ print(
66
+ "CuPy not available or not functional. Falling back to NumPy."
67
+ )
68
+ self._backend = "numpy"
37
69
  return np
38
- import numpy as np_mod
70
+ self._backend = "numpy"
71
+ return np
39
72
 
40
- return np_mod
41
-
42
- def __init_post__(self, verbose: bool = False) -> None:
43
- # This is now redundant but kept for compatibility if called
44
- self._xp = self._load_backend(self._backend, verbose)
45
- assert isinstance(self._xp, ModuleType)
46
- if verbose:
47
- print(f"Using {self._xp.__name__} backend.")
73
+ def __repr__(self) -> str:
74
+ return f"ArrayBackend(backend={self._backend!r}, module={self._xp.__name__!r})"
48
75
 
49
76
  @property
50
77
  def backend(self) -> BackendType:
@@ -70,15 +97,12 @@ class ArrayBackend:
70
97
  self._xp = old_xp
71
98
 
72
99
 
73
- # TODO: Make this configurable via environment variable or config file.
74
100
  array_backend = ArrayBackend(
75
101
  backend=(
76
102
  "cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy"
77
103
  ),
78
104
  verbose=False,
79
105
  )
80
- # Re-run initialization logic properly after backend selection
81
- array_backend.__init_post__(verbose=False)
82
106
 
83
107
 
84
108
  def use_backend(backend: BackendType) -> Generator[None, None, None]:
@@ -102,6 +126,14 @@ def _numpy_backend() -> bool:
102
126
  return array_backend.backend == "numpy"
103
127
 
104
128
 
129
+ def set_device(device_id: int) -> None:
130
+ """Select the active CUDA device for the current process (no-op on NumPy)."""
131
+ if array_backend.backend == "cupy":
132
+ import cupy as cp
133
+
134
+ cp.cuda.Device(device_id).use()
135
+
136
+
105
137
  def synchronize() -> None:
106
138
  """Wait for all kernels in all streams on current device to complete."""
107
139
  if array_backend.backend == "cupy":
@@ -109,13 +141,20 @@ def synchronize() -> None:
109
141
  import cupy as cp
110
142
 
111
143
  cp.cuda.Device().synchronize()
112
- except (ImportError, AttributeError):
144
+ except ImportError:
113
145
  pass
146
+ except AttributeError as e:
147
+ warnings.warn(
148
+ f"CuPy synchronize() failed unexpectedly, this may indicate a "
149
+ f"CuPy API mismatch: {e}",
150
+ RuntimeWarning,
151
+ stacklevel=2,
152
+ )
114
153
 
115
154
 
116
155
  def to_numpy(array: Any) -> np.ndarray:
117
156
  """Convert an array to a NumPy array."""
118
- if hasattr(array, "get"):
157
+ if get_backend(array) == "cupy":
119
158
  return array.get()
120
159
 
121
160
  return np.asarray(array)
@@ -123,25 +162,24 @@ def to_numpy(array: Any) -> np.ndarray:
123
162
 
124
163
  def to_cupy(array: Any) -> Any:
125
164
  """Convert an array to a CuPy array."""
126
- try:
127
- import cupy as cp
165
+ if not cupy_available():
166
+ raise ImportError("CuPy is not available or not functional.")
128
167
 
129
- return cp.asarray(array)
130
- except ImportError:
131
- raise ImportError("CuPy is not available.")
168
+ import array_api_compat.cupy as cp
169
+
170
+ return cp.asarray(array)
132
171
 
133
172
 
134
173
  def to_cunumpy(array: Any) -> Any:
135
174
  """Convert an array to the currently active backend."""
136
- if array_backend.backend == "cupy":
175
+ if array_backend.backend == "cupy" and cupy_available():
137
176
  return to_cupy(array)
138
177
  return to_numpy(array)
139
178
 
140
179
 
141
180
  def get_backend(array: Any) -> BackendType:
142
181
  """Return 'cupy' or 'numpy' depending on the array type."""
143
- module = getattr(type(array), "__module__", "")
144
- return "cupy" if "cupy" in module else "numpy"
182
+ return "cupy" if array_api_compat.is_cupy_array(array) else "numpy"
145
183
 
146
184
 
147
185
  def is_gpu(array: Any) -> bool:
@@ -157,7 +195,7 @@ def is_cpu(array: Any) -> bool:
157
195
  # TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
158
196
  # This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported.
159
197
  if TYPE_CHECKING:
160
- import numpy as xp
198
+ import numpy as xp # noqa: F401 - type-checker-only alias for autocompletion
161
199
  else:
162
200
  # Use module-level __getattr__ for dynamic xp (Python 3.7+)
163
201
  def __getattr__(name):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cunumpy
3
- Version: 0.1.2
3
+ Version: 0.1.4
4
4
  Summary: Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
5
5
  Author: Max
6
6
  Project-URL: Source, https://github.com/max-models/cunumpy
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
16
  Requires-Python: >=3.8
17
17
  Description-Content-Type: text/markdown
18
+ Requires-Dist: array-api-compat
18
19
  Requires-Dist: numpy
19
20
  Provides-Extra: dev
20
21
  Requires-Dist: black[jupyter]; extra == "dev"
@@ -32,6 +33,7 @@ Requires-Dist: sphinx; extra == "docs"
32
33
  Requires-Dist: sphinx-book-theme; extra == "docs"
33
34
  Provides-Extra: test
34
35
  Requires-Dist: coverage; extra == "test"
36
+ Requires-Dist: pyccel; extra == "test"
35
37
  Requires-Dist: pytest; extra == "test"
36
38
 
37
39
  # CuNumpy
@@ -52,7 +54,8 @@ export ARRAY_BACKEND=cupy
52
54
 
53
55
  ```python
54
56
  import cunumpy as xp
55
- arr = xp.array([1,2])
57
+
58
+ arr = xp.array([1, 2])
56
59
 
57
60
  print(type(arr))
58
61
  print(xp.__version__)
@@ -2,6 +2,7 @@ README.md
2
2
  pyproject.toml
3
3
  src/cunumpy/__init__.py
4
4
  src/cunumpy/__init__.pyi
5
+ src/cunumpy/kernel.py
5
6
  src/cunumpy/main.py
6
7
  src/cunumpy/py.typed
7
8
  src/cunumpy/xp.py
@@ -1,3 +1,4 @@
1
+ array-api-compat
1
2
  numpy
2
3
 
3
4
  [dev]
@@ -18,4 +19,5 @@ sphinx-book-theme
18
19
 
19
20
  [test]
20
21
  coverage
22
+ pyccel
21
23
  pytest
File without changes
File without changes
File without changes