cunumpy 0.1__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.
- {cunumpy-0.1 → cunumpy-0.1.2}/PKG-INFO +25 -13
- cunumpy-0.1.2/README.md +54 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/pyproject.toml +7 -1
- cunumpy-0.1.2/src/cunumpy/__init__.py +37 -0
- cunumpy-0.1.2/src/cunumpy/__init__.pyi +24 -0
- cunumpy-0.1.2/src/cunumpy/py.typed +0 -0
- cunumpy-0.1.2/src/cunumpy/xp.py +170 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy.egg-info/PKG-INFO +25 -13
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy.egg-info/SOURCES.txt +2 -0
- cunumpy-0.1/README.md +0 -42
- cunumpy-0.1/src/cunumpy/__init__.py +0 -9
- cunumpy-0.1/src/cunumpy/xp.py +0 -64
- {cunumpy-0.1 → cunumpy-0.1.2}/setup.cfg +0 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy/main.py +0 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy.egg-info/dependency_links.txt +0 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy.egg-info/requires.txt +0 -0
- {cunumpy-0.1 → cunumpy-0.1.2}/src/cunumpy.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cunumpy
|
|
3
|
-
Version: 0.1
|
|
3
|
+
Version: 0.1.2
|
|
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
|
|
@@ -40,18 +40,8 @@ Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cun
|
|
|
40
40
|
|
|
41
41
|
# Install
|
|
42
42
|
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
```
|
|
46
|
-
python -m venv env
|
|
47
|
-
source env/bin/activate
|
|
48
|
-
pip install --upgrade pip
|
|
49
|
-
```
|
|
50
|
-
|
|
51
|
-
Install the code and requirements with pip
|
|
52
|
-
|
|
53
|
-
```
|
|
54
|
-
pip install -e .
|
|
43
|
+
```bash
|
|
44
|
+
pip install cunumpy
|
|
55
45
|
```
|
|
56
46
|
|
|
57
47
|
Example usage:
|
|
@@ -66,6 +56,28 @@ arr = xp.array([1,2])
|
|
|
66
56
|
|
|
67
57
|
print(type(arr))
|
|
68
58
|
print(xp.__version__)
|
|
59
|
+
|
|
60
|
+
# Convert to NumPy
|
|
61
|
+
arr_np = xp.to_numpy(arr)
|
|
62
|
+
|
|
63
|
+
# Convert to active backend
|
|
64
|
+
arr_xp = xp.to_cunumpy(arr)
|
|
65
|
+
|
|
66
|
+
# Inspect backend
|
|
67
|
+
print(xp.get_backend(arr))
|
|
68
|
+
print(xp.is_gpu(arr))
|
|
69
|
+
print(xp.is_cpu(arr))
|
|
70
|
+
|
|
71
|
+
# Temporarily switch backend
|
|
72
|
+
with xp.use_backend("numpy"):
|
|
73
|
+
# This code runs on CPU even if ARRAY_BACKEND=cupy
|
|
74
|
+
arr_cpu = xp.zeros(100)
|
|
75
|
+
|
|
76
|
+
# Set backend globally
|
|
77
|
+
xp.set_backend("cupy")
|
|
78
|
+
|
|
79
|
+
# Synchronize GPU operations (no-op on CPU)
|
|
80
|
+
xp.synchronize()
|
|
69
81
|
```
|
|
70
82
|
|
|
71
83
|
# Build docs
|
cunumpy-0.1.2/README.md
ADDED
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
# CuNumpy
|
|
2
|
+
|
|
3
|
+
Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
|
|
4
|
+
|
|
5
|
+
# Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install cunumpy
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
Example usage:
|
|
12
|
+
|
|
13
|
+
```
|
|
14
|
+
export ARRAY_BACKEND=cupy
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
import cunumpy as xp
|
|
19
|
+
arr = xp.array([1,2])
|
|
20
|
+
|
|
21
|
+
print(type(arr))
|
|
22
|
+
print(xp.__version__)
|
|
23
|
+
|
|
24
|
+
# Convert to NumPy
|
|
25
|
+
arr_np = xp.to_numpy(arr)
|
|
26
|
+
|
|
27
|
+
# Convert to active backend
|
|
28
|
+
arr_xp = xp.to_cunumpy(arr)
|
|
29
|
+
|
|
30
|
+
# Inspect backend
|
|
31
|
+
print(xp.get_backend(arr))
|
|
32
|
+
print(xp.is_gpu(arr))
|
|
33
|
+
print(xp.is_cpu(arr))
|
|
34
|
+
|
|
35
|
+
# Temporarily switch backend
|
|
36
|
+
with xp.use_backend("numpy"):
|
|
37
|
+
# This code runs on CPU even if ARRAY_BACKEND=cupy
|
|
38
|
+
arr_cpu = xp.zeros(100)
|
|
39
|
+
|
|
40
|
+
# Set backend globally
|
|
41
|
+
xp.set_backend("cupy")
|
|
42
|
+
|
|
43
|
+
# Synchronize GPU operations (no-op on CPU)
|
|
44
|
+
xp.synchronize()
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
# Build docs
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
```
|
|
51
|
+
make html
|
|
52
|
+
cd ../
|
|
53
|
+
open docs/_build/html/index.html
|
|
54
|
+
```
|
|
@@ -5,7 +5,7 @@ requires = [ "setuptools", "wheel" ]
|
|
|
5
5
|
|
|
6
6
|
[project]
|
|
7
7
|
name = "cunumpy"
|
|
8
|
-
version = "0.1"
|
|
8
|
+
version = "0.1.2"
|
|
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" ]
|
|
@@ -48,3 +48,9 @@ urls."Source" = "https://github.com/max-models/cunumpy"
|
|
|
48
48
|
|
|
49
49
|
[tool.setuptools.packages.find]
|
|
50
50
|
where = [ "src" ]
|
|
51
|
+
|
|
52
|
+
[tool.setuptools.package-data]
|
|
53
|
+
cunumpy = [ "py.typed", "*.pyi" ]
|
|
54
|
+
|
|
55
|
+
[tool.isort]
|
|
56
|
+
profile = "black"
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
# cunumpy/__init__.py
|
|
2
|
+
from . import xp
|
|
3
|
+
from .xp import (
|
|
4
|
+
get_backend,
|
|
5
|
+
is_cpu,
|
|
6
|
+
is_gpu,
|
|
7
|
+
set_backend,
|
|
8
|
+
synchronize,
|
|
9
|
+
to_cunumpy,
|
|
10
|
+
to_cupy,
|
|
11
|
+
to_numpy,
|
|
12
|
+
use_backend,
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
__all__ = [
|
|
16
|
+
"xp",
|
|
17
|
+
"to_numpy",
|
|
18
|
+
"to_cupy",
|
|
19
|
+
"to_cunumpy",
|
|
20
|
+
"get_backend",
|
|
21
|
+
"is_gpu",
|
|
22
|
+
"is_cpu",
|
|
23
|
+
"use_backend",
|
|
24
|
+
"set_backend",
|
|
25
|
+
"synchronize",
|
|
26
|
+
"numpy_backend",
|
|
27
|
+
"cupy_backend",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def __getattr__(name: str):
|
|
32
|
+
"""Set cunumpy.<name> to cunumpy.xp.<name> (NumPy/CuPy)."""
|
|
33
|
+
if name == "numpy_backend":
|
|
34
|
+
return xp.numpy_backend
|
|
35
|
+
if name == "cupy_backend":
|
|
36
|
+
return xp.cupy_backend
|
|
37
|
+
return getattr(xp.xp, name)
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
# Stub file for Pylance/mypy: exposes all numpy symbols so that
|
|
2
|
+
# `import cunumpy as xp` followed by `xp.<Tab>` shows numpy completions.
|
|
3
|
+
# At runtime the real __init__.py dispatches to numpy or cupy via __getattr__.
|
|
4
|
+
from contextlib import contextmanager
|
|
5
|
+
from typing import Any, Generator
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
from numpy import *
|
|
9
|
+
|
|
10
|
+
from . import xp
|
|
11
|
+
|
|
12
|
+
def to_numpy(array: Any) -> np.ndarray: ...
|
|
13
|
+
def to_cupy(array: Any) -> Any: ...
|
|
14
|
+
def to_cunumpy(array: Any) -> Any: ...
|
|
15
|
+
def get_backend(array: Any) -> str: ...
|
|
16
|
+
def is_gpu(array: Any) -> bool: ...
|
|
17
|
+
def is_cpu(array: Any) -> bool: ...
|
|
18
|
+
@contextmanager
|
|
19
|
+
def use_backend(backend: str) -> Generator[None, None, None]: ...
|
|
20
|
+
def set_backend(backend: str) -> None: ...
|
|
21
|
+
def synchronize() -> None: ...
|
|
22
|
+
|
|
23
|
+
numpy_backend: bool
|
|
24
|
+
cupy_backend: bool
|
|
File without changes
|
|
@@ -0,0 +1,170 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from contextlib import contextmanager
|
|
3
|
+
from types import ModuleType
|
|
4
|
+
from typing import TYPE_CHECKING, Any, Generator, Literal
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
|
|
8
|
+
BackendType = Literal["numpy", "cupy"]
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class ArrayBackend:
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
backend: BackendType = "numpy",
|
|
15
|
+
verbose: bool = False,
|
|
16
|
+
) -> None:
|
|
17
|
+
assert backend.lower() in [
|
|
18
|
+
"numpy",
|
|
19
|
+
"cupy",
|
|
20
|
+
], "Array backend must be either 'numpy' or 'cupy'."
|
|
21
|
+
|
|
22
|
+
self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy"
|
|
23
|
+
self._xp: ModuleType = np # Placeholder
|
|
24
|
+
|
|
25
|
+
# Import numpy/cupy
|
|
26
|
+
self._xp = self._load_backend(self._backend, verbose)
|
|
27
|
+
|
|
28
|
+
def _load_backend(self, backend: BackendType, verbose: bool = False) -> ModuleType:
|
|
29
|
+
if backend == "cupy":
|
|
30
|
+
try:
|
|
31
|
+
import cupy as cp
|
|
32
|
+
|
|
33
|
+
return cp
|
|
34
|
+
except ImportError:
|
|
35
|
+
if verbose:
|
|
36
|
+
print("CuPy not available.")
|
|
37
|
+
return np
|
|
38
|
+
import numpy as np_mod
|
|
39
|
+
|
|
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.")
|
|
48
|
+
|
|
49
|
+
@property
|
|
50
|
+
def backend(self) -> BackendType:
|
|
51
|
+
return self._backend
|
|
52
|
+
|
|
53
|
+
@property
|
|
54
|
+
def xp(self) -> ModuleType:
|
|
55
|
+
return self._xp
|
|
56
|
+
|
|
57
|
+
@contextmanager
|
|
58
|
+
def use_backend(self, backend: BackendType) -> Generator[None, None, None]:
|
|
59
|
+
"""Temporarily change the backend."""
|
|
60
|
+
old_backend = self._backend
|
|
61
|
+
old_xp = self._xp
|
|
62
|
+
|
|
63
|
+
self._backend = backend
|
|
64
|
+
self._xp = self._load_backend(backend)
|
|
65
|
+
|
|
66
|
+
try:
|
|
67
|
+
yield
|
|
68
|
+
finally:
|
|
69
|
+
self._backend = old_backend
|
|
70
|
+
self._xp = old_xp
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
# TODO: Make this configurable via environment variable or config file.
|
|
74
|
+
array_backend = ArrayBackend(
|
|
75
|
+
backend=(
|
|
76
|
+
"cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy"
|
|
77
|
+
),
|
|
78
|
+
verbose=False,
|
|
79
|
+
)
|
|
80
|
+
# Re-run initialization logic properly after backend selection
|
|
81
|
+
array_backend.__init_post__(verbose=False)
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def use_backend(backend: BackendType) -> Generator[None, None, None]:
|
|
85
|
+
"""Temporarily change the backend."""
|
|
86
|
+
return array_backend.use_backend(backend)
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def set_backend(backend: BackendType) -> None:
|
|
90
|
+
"""Set the backend globally."""
|
|
91
|
+
array_backend._backend = backend
|
|
92
|
+
array_backend._xp = array_backend._load_backend(backend)
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def _cupy_backend() -> bool:
|
|
96
|
+
"""Check if the active global backend is CuPy."""
|
|
97
|
+
return array_backend.backend == "cupy"
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _numpy_backend() -> bool:
|
|
101
|
+
"""Check if the active global backend is NumPy."""
|
|
102
|
+
return array_backend.backend == "numpy"
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def synchronize() -> None:
|
|
106
|
+
"""Wait for all kernels in all streams on current device to complete."""
|
|
107
|
+
if array_backend.backend == "cupy":
|
|
108
|
+
try:
|
|
109
|
+
import cupy as cp
|
|
110
|
+
|
|
111
|
+
cp.cuda.Device().synchronize()
|
|
112
|
+
except (ImportError, AttributeError):
|
|
113
|
+
pass
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def to_numpy(array: Any) -> np.ndarray:
|
|
117
|
+
"""Convert an array to a NumPy array."""
|
|
118
|
+
if hasattr(array, "get"):
|
|
119
|
+
return array.get()
|
|
120
|
+
|
|
121
|
+
return np.asarray(array)
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def to_cupy(array: Any) -> Any:
|
|
125
|
+
"""Convert an array to a CuPy array."""
|
|
126
|
+
try:
|
|
127
|
+
import cupy as cp
|
|
128
|
+
|
|
129
|
+
return cp.asarray(array)
|
|
130
|
+
except ImportError:
|
|
131
|
+
raise ImportError("CuPy is not available.")
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
def to_cunumpy(array: Any) -> Any:
|
|
135
|
+
"""Convert an array to the currently active backend."""
|
|
136
|
+
if array_backend.backend == "cupy":
|
|
137
|
+
return to_cupy(array)
|
|
138
|
+
return to_numpy(array)
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def get_backend(array: Any) -> BackendType:
|
|
142
|
+
"""Return 'cupy' or 'numpy' depending on the array type."""
|
|
143
|
+
module = getattr(type(array), "__module__", "")
|
|
144
|
+
return "cupy" if "cupy" in module else "numpy"
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def is_gpu(array: Any) -> bool:
|
|
148
|
+
"""Check if the array is stored on a GPU (CuPy)."""
|
|
149
|
+
return get_backend(array) == "cupy"
|
|
150
|
+
|
|
151
|
+
|
|
152
|
+
def is_cpu(array: Any) -> bool:
|
|
153
|
+
"""Check if the array is stored on a CPU (NumPy)."""
|
|
154
|
+
return get_backend(array) == "numpy"
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
|
|
158
|
+
# This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported.
|
|
159
|
+
if TYPE_CHECKING:
|
|
160
|
+
import numpy as xp
|
|
161
|
+
else:
|
|
162
|
+
# Use module-level __getattr__ for dynamic xp (Python 3.7+)
|
|
163
|
+
def __getattr__(name):
|
|
164
|
+
if name == "xp":
|
|
165
|
+
return array_backend.xp
|
|
166
|
+
if name == "numpy_backend":
|
|
167
|
+
return _numpy_backend()
|
|
168
|
+
if name == "cupy_backend":
|
|
169
|
+
return _cupy_backend()
|
|
170
|
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: cunumpy
|
|
3
|
-
Version: 0.1
|
|
3
|
+
Version: 0.1.2
|
|
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
|
|
@@ -40,18 +40,8 @@ Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cun
|
|
|
40
40
|
|
|
41
41
|
# Install
|
|
42
42
|
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
```
|
|
46
|
-
python -m venv env
|
|
47
|
-
source env/bin/activate
|
|
48
|
-
pip install --upgrade pip
|
|
49
|
-
```
|
|
50
|
-
|
|
51
|
-
Install the code and requirements with pip
|
|
52
|
-
|
|
53
|
-
```
|
|
54
|
-
pip install -e .
|
|
43
|
+
```bash
|
|
44
|
+
pip install cunumpy
|
|
55
45
|
```
|
|
56
46
|
|
|
57
47
|
Example usage:
|
|
@@ -66,6 +56,28 @@ arr = xp.array([1,2])
|
|
|
66
56
|
|
|
67
57
|
print(type(arr))
|
|
68
58
|
print(xp.__version__)
|
|
59
|
+
|
|
60
|
+
# Convert to NumPy
|
|
61
|
+
arr_np = xp.to_numpy(arr)
|
|
62
|
+
|
|
63
|
+
# Convert to active backend
|
|
64
|
+
arr_xp = xp.to_cunumpy(arr)
|
|
65
|
+
|
|
66
|
+
# Inspect backend
|
|
67
|
+
print(xp.get_backend(arr))
|
|
68
|
+
print(xp.is_gpu(arr))
|
|
69
|
+
print(xp.is_cpu(arr))
|
|
70
|
+
|
|
71
|
+
# Temporarily switch backend
|
|
72
|
+
with xp.use_backend("numpy"):
|
|
73
|
+
# This code runs on CPU even if ARRAY_BACKEND=cupy
|
|
74
|
+
arr_cpu = xp.zeros(100)
|
|
75
|
+
|
|
76
|
+
# Set backend globally
|
|
77
|
+
xp.set_backend("cupy")
|
|
78
|
+
|
|
79
|
+
# Synchronize GPU operations (no-op on CPU)
|
|
80
|
+
xp.synchronize()
|
|
69
81
|
```
|
|
70
82
|
|
|
71
83
|
# Build docs
|
cunumpy-0.1/README.md
DELETED
|
@@ -1,42 +0,0 @@
|
|
|
1
|
-
# CuNumpy
|
|
2
|
-
|
|
3
|
-
Simple wrapper for numpy and cupy. Replace `import numpy as np` with `import cunumpy as xp`.
|
|
4
|
-
|
|
5
|
-
# Install
|
|
6
|
-
|
|
7
|
-
Create and activate python environment
|
|
8
|
-
|
|
9
|
-
```
|
|
10
|
-
python -m venv env
|
|
11
|
-
source env/bin/activate
|
|
12
|
-
pip install --upgrade pip
|
|
13
|
-
```
|
|
14
|
-
|
|
15
|
-
Install the code and requirements with pip
|
|
16
|
-
|
|
17
|
-
```
|
|
18
|
-
pip install -e .
|
|
19
|
-
```
|
|
20
|
-
|
|
21
|
-
Example usage:
|
|
22
|
-
|
|
23
|
-
```
|
|
24
|
-
export ARRAY_BACKEND=cupy
|
|
25
|
-
```
|
|
26
|
-
|
|
27
|
-
```python
|
|
28
|
-
import cunumpy as xp
|
|
29
|
-
arr = xp.array([1,2])
|
|
30
|
-
|
|
31
|
-
print(type(arr))
|
|
32
|
-
print(xp.__version__)
|
|
33
|
-
```
|
|
34
|
-
|
|
35
|
-
# Build docs
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
```
|
|
39
|
-
make html
|
|
40
|
-
cd ../
|
|
41
|
-
open docs/_build/html/index.html
|
|
42
|
-
```
|
cunumpy-0.1/src/cunumpy/xp.py
DELETED
|
@@ -1,64 +0,0 @@
|
|
|
1
|
-
import os
|
|
2
|
-
from types import ModuleType
|
|
3
|
-
from typing import TYPE_CHECKING, Literal
|
|
4
|
-
|
|
5
|
-
BackendType = Literal["numpy", "cupy"]
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
class ArrayBackend:
|
|
9
|
-
def __init__(
|
|
10
|
-
self,
|
|
11
|
-
backend: BackendType = "numpy",
|
|
12
|
-
verbose: bool = False,
|
|
13
|
-
) -> None:
|
|
14
|
-
assert backend.lower() in [
|
|
15
|
-
"numpy",
|
|
16
|
-
"cupy",
|
|
17
|
-
], "Array backend must be either 'numpy' or 'cupy'."
|
|
18
|
-
|
|
19
|
-
self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy"
|
|
20
|
-
|
|
21
|
-
# Import numpy/cupy
|
|
22
|
-
if self.backend == "cupy":
|
|
23
|
-
try:
|
|
24
|
-
import cupy as cp
|
|
25
|
-
|
|
26
|
-
self._xp = cp
|
|
27
|
-
except ImportError:
|
|
28
|
-
if verbose:
|
|
29
|
-
print("CuPy not available.")
|
|
30
|
-
self._backend = "numpy"
|
|
31
|
-
|
|
32
|
-
if self.backend == "numpy":
|
|
33
|
-
import numpy as np
|
|
34
|
-
|
|
35
|
-
self._xp = np
|
|
36
|
-
|
|
37
|
-
assert isinstance(self.xp, ModuleType)
|
|
38
|
-
|
|
39
|
-
if verbose:
|
|
40
|
-
print(f"Using {self.xp.__name__} backend.")
|
|
41
|
-
|
|
42
|
-
@property
|
|
43
|
-
def backend(self) -> BackendType:
|
|
44
|
-
return self._backend
|
|
45
|
-
|
|
46
|
-
@property
|
|
47
|
-
def xp(self) -> ModuleType:
|
|
48
|
-
return self._xp
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
# TODO: Make this configurable via environment variable or config file.
|
|
52
|
-
array_backend = ArrayBackend(
|
|
53
|
-
backend=(
|
|
54
|
-
"cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy"
|
|
55
|
-
),
|
|
56
|
-
verbose=False,
|
|
57
|
-
)
|
|
58
|
-
|
|
59
|
-
# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
|
|
60
|
-
# This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported.
|
|
61
|
-
if TYPE_CHECKING:
|
|
62
|
-
import numpy as xp
|
|
63
|
-
else:
|
|
64
|
-
xp = array_backend.xp
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|