arraybridge 0.2.0__py3-none-any.whl
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.
- arraybridge/__init__.py +37 -0
- arraybridge/conversion_helpers.py +150 -0
- arraybridge/converters.py +61 -0
- arraybridge/decorators.py +396 -0
- arraybridge/dtype_scaling.py +157 -0
- arraybridge/exceptions.py +26 -0
- arraybridge/framework_config.py +459 -0
- arraybridge/framework_ops.py +15 -0
- arraybridge/gpu_cleanup.py +149 -0
- arraybridge/oom_recovery.py +148 -0
- arraybridge/slice_processing.py +73 -0
- arraybridge/stack_utils.py +317 -0
- arraybridge/types.py +70 -0
- arraybridge/utils.py +352 -0
- arraybridge-0.2.0.dist-info/METADATA +228 -0
- arraybridge-0.2.0.dist-info/RECORD +18 -0
- arraybridge-0.2.0.dist-info/WHEEL +4 -0
- arraybridge-0.2.0.dist-info/licenses/LICENSE +21 -0
arraybridge/types.py
ADDED
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Memory type definitions for arraybridge.
|
|
3
|
+
|
|
4
|
+
This module defines the MemoryType enum and related constants for managing
|
|
5
|
+
different array/tensor frameworks.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from enum import Enum
|
|
9
|
+
from typing import Any, Callable, TypeVar
|
|
10
|
+
|
|
11
|
+
T = TypeVar('T')
|
|
12
|
+
ConversionFunc = Callable[[Any], Any]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class MemoryType(Enum):
|
|
16
|
+
"""Enum representing different array/tensor framework types."""
|
|
17
|
+
|
|
18
|
+
NUMPY = "numpy"
|
|
19
|
+
CUPY = "cupy"
|
|
20
|
+
TORCH = "torch"
|
|
21
|
+
TENSORFLOW = "tensorflow"
|
|
22
|
+
JAX = "jax"
|
|
23
|
+
PYCLESPERANTO = "pyclesperanto"
|
|
24
|
+
|
|
25
|
+
@property
|
|
26
|
+
def converter(self):
|
|
27
|
+
"""Get the converter instance for this memory type."""
|
|
28
|
+
from arraybridge.conversion_helpers import _CONVERTERS
|
|
29
|
+
return _CONVERTERS[self]
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
# Auto-generate to_X() methods on enum
|
|
33
|
+
def _add_conversion_methods():
|
|
34
|
+
"""Add to_X() conversion methods to MemoryType enum."""
|
|
35
|
+
for target_type in MemoryType:
|
|
36
|
+
method_name = f"to_{target_type.value}"
|
|
37
|
+
|
|
38
|
+
def make_method(target):
|
|
39
|
+
def method(self, data, gpu_id):
|
|
40
|
+
return getattr(self.converter, f"to_{target.value}")(data, gpu_id)
|
|
41
|
+
return method
|
|
42
|
+
|
|
43
|
+
setattr(MemoryType, method_name, make_method(target_type))
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
_add_conversion_methods()
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
# Memory type sets
|
|
50
|
+
CPU_MEMORY_TYPES: set[MemoryType] = {MemoryType.NUMPY}
|
|
51
|
+
GPU_MEMORY_TYPES: set[MemoryType] = {
|
|
52
|
+
MemoryType.CUPY,
|
|
53
|
+
MemoryType.TORCH,
|
|
54
|
+
MemoryType.TENSORFLOW,
|
|
55
|
+
MemoryType.JAX,
|
|
56
|
+
MemoryType.PYCLESPERANTO
|
|
57
|
+
}
|
|
58
|
+
SUPPORTED_MEMORY_TYPES: set[MemoryType] = CPU_MEMORY_TYPES | GPU_MEMORY_TYPES
|
|
59
|
+
|
|
60
|
+
# String value sets for validation
|
|
61
|
+
VALID_MEMORY_TYPES = {mt.value for mt in MemoryType}
|
|
62
|
+
VALID_GPU_MEMORY_TYPES = {mt.value for mt in GPU_MEMORY_TYPES}
|
|
63
|
+
|
|
64
|
+
# Memory type constants for direct access
|
|
65
|
+
MEMORY_TYPE_NUMPY = MemoryType.NUMPY.value
|
|
66
|
+
MEMORY_TYPE_CUPY = MemoryType.CUPY.value
|
|
67
|
+
MEMORY_TYPE_TORCH = MemoryType.TORCH.value
|
|
68
|
+
MEMORY_TYPE_TENSORFLOW = MemoryType.TENSORFLOW.value
|
|
69
|
+
MEMORY_TYPE_JAX = MemoryType.JAX.value
|
|
70
|
+
MEMORY_TYPE_PYCLESPERANTO = MemoryType.PYCLESPERANTO.value
|
arraybridge/utils.py
ADDED
|
@@ -0,0 +1,352 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Memory conversion utility functions for arraybridge.
|
|
3
|
+
|
|
4
|
+
This module provides utility functions for memory conversion operations,
|
|
5
|
+
supporting Clause 251 (Declarative Memory Conversion Interface) and
|
|
6
|
+
Clause 65 (Fail Loudly).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import importlib
|
|
10
|
+
import logging
|
|
11
|
+
from typing import Any, Optional
|
|
12
|
+
|
|
13
|
+
from arraybridge.types import MemoryType
|
|
14
|
+
|
|
15
|
+
from .exceptions import MemoryConversionError
|
|
16
|
+
from .framework_config import _FRAMEWORK_CONFIG
|
|
17
|
+
|
|
18
|
+
logger = logging.getLogger(__name__)
|
|
19
|
+
|
|
20
|
+
class _ModulePlaceholder:
|
|
21
|
+
"""
|
|
22
|
+
Placeholder for missing optional modules that allows attribute access
|
|
23
|
+
for type annotations while still being falsy and failing on actual use.
|
|
24
|
+
"""
|
|
25
|
+
def __init__(self, module_name: str):
|
|
26
|
+
self._module_name = module_name
|
|
27
|
+
|
|
28
|
+
def __bool__(self):
|
|
29
|
+
return False
|
|
30
|
+
|
|
31
|
+
def __getattr__(self, name):
|
|
32
|
+
# Return another placeholder for chained attribute access
|
|
33
|
+
# This allows things like cp.ndarray in type annotations to work
|
|
34
|
+
return _ModulePlaceholder(f"{self._module_name}.{name}")
|
|
35
|
+
|
|
36
|
+
def __call__(self, *args, **kwargs):
|
|
37
|
+
# If someone tries to actually call a function, fail loudly
|
|
38
|
+
raise ImportError(
|
|
39
|
+
f"Module '{self._module_name}' is not available. "
|
|
40
|
+
f"Please install the required dependency."
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
def __repr__(self):
|
|
44
|
+
return f"<ModulePlaceholder for '{self._module_name}'>"
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def optional_import(module_name: str) -> Optional[Any]:
|
|
48
|
+
"""
|
|
49
|
+
Import a module if available, otherwise return a placeholder that handles
|
|
50
|
+
attribute access gracefully for type annotations but fails on actual use.
|
|
51
|
+
|
|
52
|
+
This function allows for graceful handling of optional dependencies.
|
|
53
|
+
It can be used to import libraries that may not be installed,
|
|
54
|
+
particularly GPU-related libraries like torch, tensorflow, and cupy.
|
|
55
|
+
|
|
56
|
+
Args:
|
|
57
|
+
module_name: Name of the module to import
|
|
58
|
+
|
|
59
|
+
Returns:
|
|
60
|
+
The imported module if available, a placeholder otherwise
|
|
61
|
+
|
|
62
|
+
Example:
|
|
63
|
+
```python
|
|
64
|
+
# Import torch if available
|
|
65
|
+
torch = optional_import("torch")
|
|
66
|
+
|
|
67
|
+
# Check if torch is available before using it
|
|
68
|
+
if torch:
|
|
69
|
+
# Use torch
|
|
70
|
+
tensor = torch.tensor([1, 2, 3])
|
|
71
|
+
else:
|
|
72
|
+
# Handle the case where torch is not available
|
|
73
|
+
raise ImportError("PyTorch is required for this function")
|
|
74
|
+
```
|
|
75
|
+
"""
|
|
76
|
+
try:
|
|
77
|
+
# Use importlib.import_module which handles dotted names properly
|
|
78
|
+
return importlib.import_module(module_name)
|
|
79
|
+
except (ImportError, ModuleNotFoundError, AttributeError):
|
|
80
|
+
# Return a placeholder that handles attribute access gracefully
|
|
81
|
+
return _ModulePlaceholder(module_name)
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _ensure_module(module_name: str) -> Any:
|
|
85
|
+
"""
|
|
86
|
+
Ensure a module is imported and meets version requirements.
|
|
87
|
+
|
|
88
|
+
Args:
|
|
89
|
+
module_name: The name of the module to import
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
The imported module
|
|
93
|
+
|
|
94
|
+
Raises:
|
|
95
|
+
ImportError: If the module cannot be imported or does not meet version requirements
|
|
96
|
+
RuntimeError: If the module has known issues with specific versions
|
|
97
|
+
"""
|
|
98
|
+
try:
|
|
99
|
+
module = importlib.import_module(module_name)
|
|
100
|
+
|
|
101
|
+
# Check TensorFlow version for DLPack compatibility
|
|
102
|
+
if module_name == "tensorflow":
|
|
103
|
+
import pkg_resources
|
|
104
|
+
tf_version = pkg_resources.parse_version(module.__version__)
|
|
105
|
+
min_version = pkg_resources.parse_version("2.12.0")
|
|
106
|
+
|
|
107
|
+
if tf_version < min_version:
|
|
108
|
+
raise RuntimeError(
|
|
109
|
+
f"TensorFlow version {module.__version__} is not supported "
|
|
110
|
+
f"for DLPack operations. "
|
|
111
|
+
f"Version 2.12.0 or higher is required for stable DLPack support. "
|
|
112
|
+
f"Clause 88 (No Inferred Capabilities) violation: "
|
|
113
|
+
f"Cannot infer DLPack capability."
|
|
114
|
+
)
|
|
115
|
+
|
|
116
|
+
return module
|
|
117
|
+
except ImportError:
|
|
118
|
+
raise ImportError(
|
|
119
|
+
f"Module {module_name} is required for this operation "
|
|
120
|
+
f"but is not installed"
|
|
121
|
+
)
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def _supports_cuda_array_interface(obj: Any) -> bool:
|
|
125
|
+
"""
|
|
126
|
+
Check if an object supports the CUDA Array Interface.
|
|
127
|
+
|
|
128
|
+
Args:
|
|
129
|
+
obj: The object to check
|
|
130
|
+
|
|
131
|
+
Returns:
|
|
132
|
+
True if the object supports the CUDA Array Interface, False otherwise
|
|
133
|
+
"""
|
|
134
|
+
return hasattr(obj, "__cuda_array_interface__")
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
def _supports_dlpack(obj: Any) -> bool:
|
|
138
|
+
"""
|
|
139
|
+
Check if an object supports DLPack.
|
|
140
|
+
|
|
141
|
+
Args:
|
|
142
|
+
obj: The object to check
|
|
143
|
+
|
|
144
|
+
Returns:
|
|
145
|
+
True if the object supports DLPack, False otherwise
|
|
146
|
+
|
|
147
|
+
Note:
|
|
148
|
+
For TensorFlow tensors, this function enforces Clause 88 (No Inferred Capabilities)
|
|
149
|
+
by explicitly checking:
|
|
150
|
+
1. TensorFlow version must be 2.12+ for stable DLPack support
|
|
151
|
+
2. Tensor must be on GPU (CPU tensors might succeed even without proper DLPack support)
|
|
152
|
+
3. tf.experimental.dlpack module must exist
|
|
153
|
+
"""
|
|
154
|
+
# Check for PyTorch, CuPy, or JAX DLPack support
|
|
155
|
+
# PyTorch: __dlpack__ method, CuPy: toDlpack method, JAX: __dlpack__ method
|
|
156
|
+
if hasattr(obj, "toDlpack") or hasattr(obj, "to_dlpack") or hasattr(obj, "__dlpack__"):
|
|
157
|
+
# Special handling for TensorFlow to enforce Clause 88
|
|
158
|
+
if 'tensorflow' in str(type(obj)):
|
|
159
|
+
try:
|
|
160
|
+
import tensorflow as tf
|
|
161
|
+
|
|
162
|
+
# Check TensorFlow version - DLPack is only stable in TF 2.12+
|
|
163
|
+
tf_version = tf.__version__
|
|
164
|
+
major, minor = map(int, tf_version.split('.')[:2])
|
|
165
|
+
|
|
166
|
+
if major < 2 or (major == 2 and minor < 12):
|
|
167
|
+
# Explicitly fail for TF < 2.12 to prevent silent fallbacks
|
|
168
|
+
raise RuntimeError(
|
|
169
|
+
f"TensorFlow version {tf_version} does not support "
|
|
170
|
+
f"stable DLPack operations. "
|
|
171
|
+
f"Version 2.12.0 or higher is required. "
|
|
172
|
+
f"Clause 88 violation: Cannot infer DLPack capability."
|
|
173
|
+
)
|
|
174
|
+
|
|
175
|
+
# Check if tensor is on GPU - CPU tensors might succeed
|
|
176
|
+
# even without proper DLPack support
|
|
177
|
+
device_str = obj.device.lower()
|
|
178
|
+
if "gpu" not in device_str:
|
|
179
|
+
# Explicitly fail for CPU tensors to prevent deceptive behavior
|
|
180
|
+
raise RuntimeError(
|
|
181
|
+
"TensorFlow tensor on CPU cannot use DLPack operations reliably. "
|
|
182
|
+
"Only GPU tensors are supported for DLPack operations. "
|
|
183
|
+
"Clause 88 violation: Cannot infer GPU capability."
|
|
184
|
+
)
|
|
185
|
+
|
|
186
|
+
# Check if experimental.dlpack module exists
|
|
187
|
+
if not hasattr(tf.experimental, "dlpack"):
|
|
188
|
+
raise RuntimeError(
|
|
189
|
+
"TensorFlow installation missing experimental.dlpack module. "
|
|
190
|
+
"Clause 88 violation: Cannot infer DLPack capability."
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
return True
|
|
194
|
+
except (ImportError, AttributeError) as e:
|
|
195
|
+
# Re-raise with more specific error message
|
|
196
|
+
raise RuntimeError(
|
|
197
|
+
f"TensorFlow DLPack support check failed: {str(e)}. "
|
|
198
|
+
f"Clause 88 violation: Cannot infer DLPack capability."
|
|
199
|
+
) from e
|
|
200
|
+
|
|
201
|
+
# For non-TensorFlow types, return True if they have DLPack methods
|
|
202
|
+
return True
|
|
203
|
+
|
|
204
|
+
return False
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
# NOTE: Device operations now defined in framework_config.py
|
|
208
|
+
# This eliminates the scattered _DEVICE_OPS dict
|
|
209
|
+
|
|
210
|
+
|
|
211
|
+
def _get_device_id(data: Any, memory_type: str) -> Optional[int]:
|
|
212
|
+
"""
|
|
213
|
+
Get the GPU device ID from a data object using framework config.
|
|
214
|
+
|
|
215
|
+
Args:
|
|
216
|
+
data: The data object
|
|
217
|
+
memory_type: The memory type
|
|
218
|
+
|
|
219
|
+
Returns:
|
|
220
|
+
The GPU device ID or None if not applicable
|
|
221
|
+
|
|
222
|
+
Raises:
|
|
223
|
+
MemoryConversionError: If the device ID cannot be determined for a GPU memory type
|
|
224
|
+
"""
|
|
225
|
+
# Convert string to enum
|
|
226
|
+
mem_type = MemoryType(memory_type)
|
|
227
|
+
config = _FRAMEWORK_CONFIG[mem_type]
|
|
228
|
+
get_id_handler = config['get_device_id']
|
|
229
|
+
|
|
230
|
+
# Check if it's a callable handler (pyclesperanto)
|
|
231
|
+
if callable(get_id_handler):
|
|
232
|
+
mod = _ensure_module(mem_type.value)
|
|
233
|
+
return get_id_handler(data, mod)
|
|
234
|
+
|
|
235
|
+
# Check if it's None (CPU)
|
|
236
|
+
if get_id_handler is None:
|
|
237
|
+
return None
|
|
238
|
+
|
|
239
|
+
# It's an eval expression
|
|
240
|
+
try:
|
|
241
|
+
mod = _ensure_module(mem_type.value) # noqa: F841 (used in eval)
|
|
242
|
+
return eval(get_id_handler)
|
|
243
|
+
except (AttributeError, Exception) as e:
|
|
244
|
+
logger.warning(f"Failed to get device ID for {mem_type.value} array: {e}")
|
|
245
|
+
# Try fallback if available
|
|
246
|
+
if 'get_device_id_fallback' in config:
|
|
247
|
+
return eval(config['get_device_id_fallback'])
|
|
248
|
+
|
|
249
|
+
|
|
250
|
+
def _set_device(memory_type: str, device_id: int) -> None:
|
|
251
|
+
"""
|
|
252
|
+
Set the current device for a specific memory type using framework config.
|
|
253
|
+
|
|
254
|
+
Args:
|
|
255
|
+
memory_type: The memory type
|
|
256
|
+
device_id: The GPU device ID
|
|
257
|
+
|
|
258
|
+
Raises:
|
|
259
|
+
MemoryConversionError: If the device cannot be set
|
|
260
|
+
"""
|
|
261
|
+
# Convert string to enum
|
|
262
|
+
mem_type = MemoryType(memory_type)
|
|
263
|
+
config = _FRAMEWORK_CONFIG[mem_type]
|
|
264
|
+
set_device_handler = config['set_device']
|
|
265
|
+
|
|
266
|
+
# Check if it's a callable handler (pyclesperanto)
|
|
267
|
+
if callable(set_device_handler):
|
|
268
|
+
try:
|
|
269
|
+
mod = _ensure_module(mem_type.value)
|
|
270
|
+
set_device_handler(device_id, mod)
|
|
271
|
+
except Exception as e:
|
|
272
|
+
raise MemoryConversionError(
|
|
273
|
+
source_type=memory_type,
|
|
274
|
+
target_type=memory_type,
|
|
275
|
+
method="device_selection",
|
|
276
|
+
reason=f"Failed to set {mem_type.value} device to {device_id}: {e}"
|
|
277
|
+
) from e
|
|
278
|
+
return
|
|
279
|
+
|
|
280
|
+
# Check if it's None (frameworks that don't need global device setting)
|
|
281
|
+
if set_device_handler is None:
|
|
282
|
+
return
|
|
283
|
+
|
|
284
|
+
# It's an eval expression
|
|
285
|
+
try:
|
|
286
|
+
mod = _ensure_module(mem_type.value) # noqa: F841 (used in eval)
|
|
287
|
+
eval(set_device_handler.format(mod='mod'))
|
|
288
|
+
except Exception as e:
|
|
289
|
+
raise MemoryConversionError(
|
|
290
|
+
source_type=memory_type,
|
|
291
|
+
target_type=memory_type,
|
|
292
|
+
method="device_selection",
|
|
293
|
+
reason=f"Failed to set {mem_type.value} device to {device_id}: {e}"
|
|
294
|
+
) from e
|
|
295
|
+
|
|
296
|
+
|
|
297
|
+
def _move_to_device(data: Any, memory_type: str, device_id: int) -> Any:
|
|
298
|
+
"""
|
|
299
|
+
Move data to a specific GPU device using framework config.
|
|
300
|
+
|
|
301
|
+
Args:
|
|
302
|
+
data: The data to move
|
|
303
|
+
memory_type: The memory type
|
|
304
|
+
device_id: The target GPU device ID
|
|
305
|
+
|
|
306
|
+
Returns:
|
|
307
|
+
The data on the target device
|
|
308
|
+
|
|
309
|
+
Raises:
|
|
310
|
+
MemoryConversionError: If the data cannot be moved to the specified device
|
|
311
|
+
"""
|
|
312
|
+
# Convert string to enum
|
|
313
|
+
mem_type = MemoryType(memory_type)
|
|
314
|
+
config = _FRAMEWORK_CONFIG[mem_type]
|
|
315
|
+
move_handler = config['move_to_device']
|
|
316
|
+
|
|
317
|
+
# Check if it's a callable handler (pyclesperanto)
|
|
318
|
+
if callable(move_handler):
|
|
319
|
+
try:
|
|
320
|
+
mod = _ensure_module(mem_type.value)
|
|
321
|
+
return move_handler(data, device_id, mod, memory_type)
|
|
322
|
+
except Exception as e:
|
|
323
|
+
raise MemoryConversionError(
|
|
324
|
+
source_type=memory_type,
|
|
325
|
+
target_type=memory_type,
|
|
326
|
+
method="device_movement",
|
|
327
|
+
reason=f"Failed to move {mem_type.value} array to device {device_id}: {e}"
|
|
328
|
+
) from e
|
|
329
|
+
|
|
330
|
+
# Check if it's None (CPU memory types)
|
|
331
|
+
if move_handler is None:
|
|
332
|
+
return data
|
|
333
|
+
|
|
334
|
+
# It's an eval expression
|
|
335
|
+
try:
|
|
336
|
+
mod = _ensure_module(mem_type.value) # noqa: F841 (used in eval)
|
|
337
|
+
|
|
338
|
+
# Handle context managers (CuPy, TensorFlow)
|
|
339
|
+
if 'move_context' in config and config['move_context']:
|
|
340
|
+
context_expr = config['move_context'].format(mod='mod')
|
|
341
|
+
context = eval(context_expr)
|
|
342
|
+
with context:
|
|
343
|
+
return eval(move_handler.format(mod='mod'))
|
|
344
|
+
else:
|
|
345
|
+
return eval(move_handler.format(mod='mod'))
|
|
346
|
+
except Exception as e:
|
|
347
|
+
raise MemoryConversionError(
|
|
348
|
+
source_type=memory_type,
|
|
349
|
+
target_type=memory_type,
|
|
350
|
+
method="device_movement",
|
|
351
|
+
reason=f"Failed to move {mem_type.value} array to device {device_id}: {e}"
|
|
352
|
+
) from e
|
|
@@ -0,0 +1,228 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: arraybridge
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion
|
|
5
|
+
Project-URL: Homepage, https://github.com/trissim/arraybridge
|
|
6
|
+
Project-URL: Documentation, https://arraybridge.readthedocs.io
|
|
7
|
+
Project-URL: Repository, https://github.com/trissim/arraybridge
|
|
8
|
+
Project-URL: Issues, https://github.com/trissim/arraybridge/issues
|
|
9
|
+
Author-email: Tristan Simas <tristan.simas@mail.mcgill.ca>
|
|
10
|
+
License: MIT
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Keywords: array,conversion,cupy,gpu,jax,numpy,pytorch,tensor,tensorflow
|
|
13
|
+
Classifier: Development Status :: 4 - Beta
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Programming Language :: Python :: 3
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering
|
|
23
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
24
|
+
Requires-Python: >=3.9
|
|
25
|
+
Requires-Dist: numpy>=1.20
|
|
26
|
+
Provides-Extra: all
|
|
27
|
+
Requires-Dist: cupy>=10.0; extra == 'all'
|
|
28
|
+
Requires-Dist: jax>=0.3; extra == 'all'
|
|
29
|
+
Requires-Dist: jaxlib>=0.3; extra == 'all'
|
|
30
|
+
Requires-Dist: pyclesperanto>=0.10; extra == 'all'
|
|
31
|
+
Requires-Dist: tensorflow>=2.8; extra == 'all'
|
|
32
|
+
Requires-Dist: torch>=1.10; extra == 'all'
|
|
33
|
+
Provides-Extra: cupy
|
|
34
|
+
Requires-Dist: cupy>=10.0; extra == 'cupy'
|
|
35
|
+
Provides-Extra: dev
|
|
36
|
+
Requires-Dist: black>=23.0; extra == 'dev'
|
|
37
|
+
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
38
|
+
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
|
|
39
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
40
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
41
|
+
Provides-Extra: docs
|
|
42
|
+
Requires-Dist: sphinx-autodoc-typehints>=1.25; extra == 'docs'
|
|
43
|
+
Requires-Dist: sphinx-rtd-theme>=2.0; extra == 'docs'
|
|
44
|
+
Requires-Dist: sphinx>=7.0; extra == 'docs'
|
|
45
|
+
Provides-Extra: jax
|
|
46
|
+
Requires-Dist: jax>=0.3; extra == 'jax'
|
|
47
|
+
Requires-Dist: jaxlib>=0.3; extra == 'jax'
|
|
48
|
+
Provides-Extra: pyclesperanto
|
|
49
|
+
Requires-Dist: pyclesperanto>=0.10; extra == 'pyclesperanto'
|
|
50
|
+
Provides-Extra: tensorflow
|
|
51
|
+
Requires-Dist: tensorflow>=2.8; extra == 'tensorflow'
|
|
52
|
+
Provides-Extra: torch
|
|
53
|
+
Requires-Dist: torch>=1.10; extra == 'torch'
|
|
54
|
+
Description-Content-Type: text/markdown
|
|
55
|
+
|
|
56
|
+
# arraybridge
|
|
57
|
+
|
|
58
|
+
**Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto**
|
|
59
|
+
|
|
60
|
+
[](https://badge.fury.io/py/arraybridge)
|
|
61
|
+
[](https://www.python.org/downloads/)
|
|
62
|
+
[](https://opensource.org/licenses/MIT)
|
|
63
|
+
|
|
64
|
+
## Features
|
|
65
|
+
|
|
66
|
+
- **Unified API**: Single interface for 6 array/tensor frameworks
|
|
67
|
+
- **Automatic Conversion**: DLPack + NumPy fallback with automatic path selection
|
|
68
|
+
- **Declarative Decorators**: `@numpy`, `@torch`, `@cupy` for memory type declarations
|
|
69
|
+
- **Device Management**: Thread-local GPU contexts and automatic stream management
|
|
70
|
+
- **OOM Recovery**: Automatic out-of-memory detection and cache clearing
|
|
71
|
+
- **Dtype Preservation**: Automatic dtype preservation across conversions
|
|
72
|
+
- **Zero Dependencies**: Only requires NumPy (framework dependencies are optional)
|
|
73
|
+
|
|
74
|
+
## Quick Start
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from arraybridge import convert_memory, detect_memory_type
|
|
78
|
+
import numpy as np
|
|
79
|
+
|
|
80
|
+
# Create NumPy array
|
|
81
|
+
data = np.array([[1, 2], [3, 4]])
|
|
82
|
+
|
|
83
|
+
# Convert to PyTorch (if installed)
|
|
84
|
+
torch_data = convert_memory(data, source_type='numpy', target_type='torch', gpu_id=0)
|
|
85
|
+
|
|
86
|
+
# Detect memory type
|
|
87
|
+
mem_type = detect_memory_type(torch_data) # 'torch'
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Declarative Decorators
|
|
91
|
+
|
|
92
|
+
```python
|
|
93
|
+
from arraybridge import numpy, torch, cupy
|
|
94
|
+
|
|
95
|
+
@torch(input_type='numpy', output_type='torch', oom_recovery=True)
|
|
96
|
+
def my_gpu_function(data):
|
|
97
|
+
"""Automatically converts input from NumPy to PyTorch."""
|
|
98
|
+
return data * 2
|
|
99
|
+
|
|
100
|
+
# Use with NumPy input
|
|
101
|
+
result = my_gpu_function(np.array([1, 2, 3])) # Returns PyTorch tensor
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
## Installation
|
|
105
|
+
|
|
106
|
+
```bash
|
|
107
|
+
# Base installation (NumPy only)
|
|
108
|
+
pip install arraybridge
|
|
109
|
+
|
|
110
|
+
# With specific frameworks
|
|
111
|
+
pip install arraybridge[torch]
|
|
112
|
+
pip install arraybridge[cupy]
|
|
113
|
+
pip install arraybridge[tensorflow]
|
|
114
|
+
pip install arraybridge[jax]
|
|
115
|
+
pip install arraybridge[pyclesperanto]
|
|
116
|
+
|
|
117
|
+
# With all frameworks
|
|
118
|
+
pip install arraybridge[all]
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
## Supported Frameworks
|
|
122
|
+
|
|
123
|
+
| Framework | CPU | GPU | DLPack | Notes |
|
|
124
|
+
|-----------|-----|-----|--------|-------|
|
|
125
|
+
| NumPy | ✅ | ❌ | ❌ | Base framework |
|
|
126
|
+
| CuPy | ❌ | ✅ | ✅ | CUDA arrays |
|
|
127
|
+
| PyTorch | ✅ | ✅ | ✅ | Tensors |
|
|
128
|
+
| TensorFlow | ✅ | ✅ | ✅ | Tensors |
|
|
129
|
+
| JAX | ✅ | ✅ | ✅ | Arrays |
|
|
130
|
+
| pyclesperanto | ❌ | ✅ | ❌ | OpenCL arrays |
|
|
131
|
+
|
|
132
|
+
## Why arraybridge?
|
|
133
|
+
|
|
134
|
+
**Before** (Manual conversion hell):
|
|
135
|
+
```python
|
|
136
|
+
import numpy as np
|
|
137
|
+
import torch
|
|
138
|
+
import cupy as cp
|
|
139
|
+
|
|
140
|
+
def process_data(data, target='torch'):
|
|
141
|
+
if target == 'torch':
|
|
142
|
+
if isinstance(data, np.ndarray):
|
|
143
|
+
return torch.from_numpy(data).cuda()
|
|
144
|
+
elif isinstance(data, cp.ndarray):
|
|
145
|
+
return torch.as_tensor(data, device='cuda')
|
|
146
|
+
elif target == 'cupy':
|
|
147
|
+
if isinstance(data, np.ndarray):
|
|
148
|
+
return cp.asarray(data)
|
|
149
|
+
elif hasattr(data, '__cuda_array_interface__'):
|
|
150
|
+
return cp.asarray(data)
|
|
151
|
+
# ... 30 more lines of if/elif ...
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
**After** (arraybridge):
|
|
155
|
+
```python
|
|
156
|
+
from arraybridge import convert_memory, detect_memory_type
|
|
157
|
+
|
|
158
|
+
def process_data(data, target='torch'):
|
|
159
|
+
source = detect_memory_type(data)
|
|
160
|
+
return convert_memory(data, source_type=source, target_type=target, gpu_id=0)
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
## Advanced Features
|
|
164
|
+
|
|
165
|
+
### Thread-Local GPU Streams
|
|
166
|
+
|
|
167
|
+
```python
|
|
168
|
+
from arraybridge import torch
|
|
169
|
+
|
|
170
|
+
@torch(oom_recovery=True)
|
|
171
|
+
def parallel_processing(data):
|
|
172
|
+
# Automatically uses thread-local CUDA stream
|
|
173
|
+
# Enables true parallelization across threads
|
|
174
|
+
return data * 2
|
|
175
|
+
```
|
|
176
|
+
|
|
177
|
+
### OOM Recovery
|
|
178
|
+
|
|
179
|
+
```python
|
|
180
|
+
from arraybridge import cupy
|
|
181
|
+
|
|
182
|
+
@cupy(oom_recovery=True)
|
|
183
|
+
def memory_intensive_operation(data):
|
|
184
|
+
# Automatically catches OOM errors
|
|
185
|
+
# Clears GPU cache and retries
|
|
186
|
+
return data @ data.T
|
|
187
|
+
```
|
|
188
|
+
|
|
189
|
+
### Stack Utilities
|
|
190
|
+
|
|
191
|
+
```python
|
|
192
|
+
from arraybridge import stack_slices, unstack_slices
|
|
193
|
+
|
|
194
|
+
# Stack 2D slices into 3D array
|
|
195
|
+
slices_2d = [np.random.rand(100, 100) for _ in range(50)]
|
|
196
|
+
volume_3d = stack_slices(slices_2d, memory_type='torch', gpu_id=0)
|
|
197
|
+
|
|
198
|
+
# Unstack 3D array into 2D slices
|
|
199
|
+
slices_back = unstack_slices(volume_3d, memory_type='torch')
|
|
200
|
+
```
|
|
201
|
+
|
|
202
|
+
## Documentation
|
|
203
|
+
|
|
204
|
+
Full documentation available at [arraybridge.readthedocs.io](https://arraybridge.readthedocs.io)
|
|
205
|
+
|
|
206
|
+
## Performance
|
|
207
|
+
|
|
208
|
+
arraybridge uses DLPack for zero-copy conversions when possible:
|
|
209
|
+
|
|
210
|
+
| Conversion | Method | Speed |
|
|
211
|
+
|------------|--------|-------|
|
|
212
|
+
| NumPy → PyTorch | `torch.from_numpy()` | Zero-copy |
|
|
213
|
+
| PyTorch → CuPy | DLPack | Zero-copy |
|
|
214
|
+
| CuPy → JAX | DLPack | Zero-copy |
|
|
215
|
+
| NumPy → CuPy | Copy | Fast |
|
|
216
|
+
| PyTorch → NumPy | `.numpy()` | Zero-copy (CPU) |
|
|
217
|
+
|
|
218
|
+
## License
|
|
219
|
+
|
|
220
|
+
MIT License - see LICENSE file for details
|
|
221
|
+
|
|
222
|
+
## Contributing
|
|
223
|
+
|
|
224
|
+
Contributions welcome! Please see CONTRIBUTING.md for guidelines.
|
|
225
|
+
|
|
226
|
+
## Credits
|
|
227
|
+
|
|
228
|
+
Developed by Tristan Simas as part of the OpenHCS project.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
arraybridge/__init__.py,sha256=oOdSZQZ1Jv5TeN-f9tdn0Ig3Yy5Z_W795TcI3l10BWM,998
|
|
2
|
+
arraybridge/conversion_helpers.py,sha256=VJZybT2EjkQ42u8wrnrQlgcDKFdc4EVQ5BTVmZcEojA,5556
|
|
3
|
+
arraybridge/converters.py,sha256=5tvw--rZTu2DR8rbHl7RIuYW29EUhA9ScGkUnIhfQco,1927
|
|
4
|
+
arraybridge/decorators.py,sha256=zH_hLIXkuoSH0yiG7uj7YtIRE82FHwGw9e-g3xUGIZU,14269
|
|
5
|
+
arraybridge/dtype_scaling.py,sha256=I2ci6VeSzB4V4_gLtqG-nz2NkHnsfokr1aR-tQg-kak,5565
|
|
6
|
+
arraybridge/exceptions.py,sha256=fhej5ZS39QcBFWzNeNYjoqeLVnOgwwp-cWIme25jyuo,752
|
|
7
|
+
arraybridge/framework_config.py,sha256=h1k5MZhENhQoy-MTCMSrTecX3u9QDnTqwVaSHRzEQHA,17948
|
|
8
|
+
arraybridge/framework_ops.py,sha256=Na5e_SkM6l3ZcmCdXowavdUkprYrwhc45hKFDhmd3X8,462
|
|
9
|
+
arraybridge/gpu_cleanup.py,sha256=uvHD-ATYBg5LCeTKHdTh7VseZ8UecjeTnpvRuF_LI3g,5118
|
|
10
|
+
arraybridge/oom_recovery.py,sha256=KP6Q_1itOMi2qPCvm1WcA7BksRfNwt68RxryqpItl-U,4591
|
|
11
|
+
arraybridge/slice_processing.py,sha256=jVJjBp5AzPpQPjnUFnCN1xemRDn6-qY6HfKopkUeklo,2693
|
|
12
|
+
arraybridge/stack_utils.py,sha256=jwHkF7LeuP9rMGSwjWg0HtQoHv8LdFHgpMtysRxeLN8,11211
|
|
13
|
+
arraybridge/types.py,sha256=7m3N8kA0Y9m04HyG3VebglT0CP8SqHJBRzYofN7AX5A,1999
|
|
14
|
+
arraybridge/utils.py,sha256=X1_weXzXLKm3wZd9BKHDXEnzFgaDEovIy_qwj9OfrcA,12313
|
|
15
|
+
arraybridge-0.2.0.dist-info/METADATA,sha256=P5FAPjx4x3FUQdrBSms_FpfEA3O4ao5PjfKUPEPw4ds,7155
|
|
16
|
+
arraybridge-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
17
|
+
arraybridge-0.2.0.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
|
|
18
|
+
arraybridge-0.2.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Tristan Simas
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|