cunumpy 0.1.2__tar.gz → 0.1.3__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.3
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
@@ -52,7 +52,8 @@ export ARRAY_BACKEND=cupy
52
52
 
53
53
  ```python
54
54
  import cunumpy as xp
55
- arr = xp.array([1,2])
55
+
56
+ arr = xp.array([1, 2])
56
57
 
57
58
  print(type(arr))
58
59
  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.3"
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" ]
@@ -1,10 +1,14 @@
1
1
  # cunumpy/__init__.py
2
+ from importlib.metadata import PackageNotFoundError, version
3
+
2
4
  from . import xp
3
5
  from .xp import (
6
+ cupy_available,
4
7
  get_backend,
5
8
  is_cpu,
6
9
  is_gpu,
7
10
  set_backend,
11
+ set_device,
8
12
  synchronize,
9
13
  to_cunumpy,
10
14
  to_cupy,
@@ -12,19 +16,27 @@ from .xp import (
12
16
  use_backend,
13
17
  )
14
18
 
19
+ try:
20
+ __version__ = version("cunumpy")
21
+ except PackageNotFoundError:
22
+ __version__ = "0.0.0+unknown"
23
+
15
24
  __all__ = [
16
- "xp",
17
- "to_numpy",
18
- "to_cupy",
19
- "to_cunumpy",
25
+ "__version__",
26
+ "cupy_available",
27
+ "cupy_backend",
20
28
  "get_backend",
21
- "is_gpu",
22
29
  "is_cpu",
23
- "use_backend",
30
+ "is_gpu",
31
+ "numpy_backend",
24
32
  "set_backend",
33
+ "set_device",
25
34
  "synchronize",
26
- "numpy_backend",
27
- "cupy_backend",
35
+ "to_cunumpy",
36
+ "to_cupy",
37
+ "to_numpy",
38
+ "use_backend",
39
+ "xp",
28
40
  ]
29
41
 
30
42
 
@@ -7,18 +7,21 @@ 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
11
 
12
12
  def to_numpy(array: Any) -> np.ndarray: ...
13
13
  def to_cupy(array: Any) -> Any: ...
14
14
  def to_cunumpy(array: Any) -> Any: ...
15
+ def cupy_available() -> bool: ...
15
16
  def get_backend(array: Any) -> str: ...
16
17
  def is_gpu(array: Any) -> bool: ...
17
18
  def is_cpu(array: Any) -> bool: ...
18
19
  @contextmanager
19
- def use_backend(backend: str) -> Generator[None, None, None]: ...
20
+ def use_backend(backend: str) -> Generator[None]: ...
20
21
  def set_backend(backend: str) -> None: ...
22
+ def set_device(device_id: int) -> None: ...
21
23
  def synchronize() -> None: ...
22
24
 
23
25
  numpy_backend: bool
24
26
  cupy_backend: bool
27
+ __version__: str
@@ -1,4 +1,5 @@
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
@@ -8,16 +9,42 @@ import numpy as np
8
9
  BackendType = Literal["numpy", "cupy"]
9
10
 
10
11
 
12
+ _CUPY_AVAILABLE_CACHE = None
13
+
14
+
15
+ def cupy_available() -> bool:
16
+ """Check if CuPy is available and functional."""
17
+ global _CUPY_AVAILABLE_CACHE
18
+ if _CUPY_AVAILABLE_CACHE is not None:
19
+ return _CUPY_AVAILABLE_CACHE
20
+
21
+ try:
22
+ import cupy as cp
23
+
24
+ # Check if a GPU is available
25
+ _CUPY_AVAILABLE_CACHE = cp.is_available()
26
+ return _CUPY_AVAILABLE_CACHE
27
+ except Exception: # noqa: BLE001 - tolerate any driver/runtime failure
28
+ _CUPY_AVAILABLE_CACHE = False
29
+ return False
30
+
31
+
11
32
  class ArrayBackend:
33
+ """Holds the process-wide active backend (NumPy or CuPy).
34
+
35
+ Not thread-safe: `set_backend`/`use_backend` mutate this single shared
36
+ instance in place, so concurrent code (threads, async tasks) switching
37
+ backends independently will race. Safe for the typical single-threaded
38
+ script/notebook usage this library targets.
39
+ """
40
+
12
41
  def __init__(
13
42
  self,
14
43
  backend: BackendType = "numpy",
15
44
  verbose: bool = False,
16
45
  ) -> None:
17
- assert backend.lower() in [
18
- "numpy",
19
- "cupy",
20
- ], "Array backend must be either 'numpy' or 'cupy'."
46
+ if backend.lower() not in ("numpy", "cupy"):
47
+ raise ValueError("Array backend must be either 'numpy' or 'cupy'.")
21
48
 
22
49
  self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy"
23
50
  self._xp: ModuleType = np # Placeholder
@@ -27,24 +54,25 @@ class ArrayBackend:
27
54
 
28
55
  def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleType:
29
56
  if backend == "cupy":
30
- try:
57
+ if cupy_available():
31
58
  import cupy as cp
32
59
 
60
+ self._backend = "cupy"
33
61
  return cp
34
- except ImportError:
62
+ else:
35
63
  if verbose:
36
- print("CuPy not available.")
64
+ print(
65
+ "CuPy not available or not functional. Falling back to NumPy."
66
+ )
67
+ self._backend = "numpy"
37
68
  return np
38
69
  import numpy as np_mod
39
70
 
71
+ self._backend = "numpy"
40
72
  return np_mod
41
73
 
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.")
74
+ def __repr__(self) -> str:
75
+ return f"ArrayBackend(backend={self._backend!r}, module={self._xp.__name__!r})"
48
76
 
49
77
  @property
50
78
  def backend(self) -> BackendType:
@@ -70,15 +98,12 @@ class ArrayBackend:
70
98
  self._xp = old_xp
71
99
 
72
100
 
73
- # TODO: Make this configurable via environment variable or config file.
74
101
  array_backend = ArrayBackend(
75
102
  backend=(
76
103
  "cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy"
77
104
  ),
78
105
  verbose=False,
79
106
  )
80
- # Re-run initialization logic properly after backend selection
81
- array_backend.__init_post__(verbose=False)
82
107
 
83
108
 
84
109
  def use_backend(backend: BackendType) -> Generator[None, None, None]:
@@ -102,6 +127,14 @@ def _numpy_backend() -> bool:
102
127
  return array_backend.backend == "numpy"
103
128
 
104
129
 
130
+ def set_device(device_id: int) -> None:
131
+ """Select the active CUDA device for the current process (no-op on NumPy)."""
132
+ if array_backend.backend == "cupy":
133
+ import cupy as cp
134
+
135
+ cp.cuda.Device(device_id).use()
136
+
137
+
105
138
  def synchronize() -> None:
106
139
  """Wait for all kernels in all streams on current device to complete."""
107
140
  if array_backend.backend == "cupy":
@@ -109,13 +142,20 @@ def synchronize() -> None:
109
142
  import cupy as cp
110
143
 
111
144
  cp.cuda.Device().synchronize()
112
- except (ImportError, AttributeError):
145
+ except ImportError:
113
146
  pass
147
+ except AttributeError as e:
148
+ warnings.warn(
149
+ f"CuPy synchronize() failed unexpectedly, this may indicate a "
150
+ f"CuPy API mismatch: {e}",
151
+ RuntimeWarning,
152
+ stacklevel=2,
153
+ )
114
154
 
115
155
 
116
156
  def to_numpy(array: Any) -> np.ndarray:
117
157
  """Convert an array to a NumPy array."""
118
- if hasattr(array, "get"):
158
+ if get_backend(array) == "cupy":
119
159
  return array.get()
120
160
 
121
161
  return np.asarray(array)
@@ -123,17 +163,17 @@ def to_numpy(array: Any) -> np.ndarray:
123
163
 
124
164
  def to_cupy(array: Any) -> Any:
125
165
  """Convert an array to a CuPy array."""
126
- try:
127
- import cupy as cp
166
+ if not cupy_available():
167
+ raise ImportError("CuPy is not available or not functional.")
128
168
 
129
- return cp.asarray(array)
130
- except ImportError:
131
- raise ImportError("CuPy is not available.")
169
+ import cupy as cp
170
+
171
+ return cp.asarray(array)
132
172
 
133
173
 
134
174
  def to_cunumpy(array: Any) -> Any:
135
175
  """Convert an array to the currently active backend."""
136
- if array_backend.backend == "cupy":
176
+ if array_backend.backend == "cupy" and cupy_available():
137
177
  return to_cupy(array)
138
178
  return to_numpy(array)
139
179
 
@@ -157,7 +197,7 @@ def is_cpu(array: Any) -> bool:
157
197
  # TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
158
198
  # This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported.
159
199
  if TYPE_CHECKING:
160
- import numpy as xp
200
+ import numpy as xp # noqa: F401 - type-checker-only alias for autocompletion
161
201
  else:
162
202
  # Use module-level __getattr__ for dynamic xp (Python 3.7+)
163
203
  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.3
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
@@ -52,7 +52,8 @@ export ARRAY_BACKEND=cupy
52
52
 
53
53
  ```python
54
54
  import cunumpy as xp
55
- arr = xp.array([1,2])
55
+
56
+ arr = xp.array([1, 2])
56
57
 
57
58
  print(type(arr))
58
59
  print(xp.__version__)
File without changes
File without changes
File without changes