arraybridge 0.2.0__py3-none-any.whl → 0.2.8__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 +35 -2
- arraybridge/converters.py +17 -7
- arraybridge/converters_registry.py +196 -0
- arraybridge/decorators.py +92 -88
- arraybridge/dtype_scaling.py +52 -30
- arraybridge/framework_config.py +233 -250
- arraybridge/framework_ops.py +0 -1
- arraybridge/gpu_cleanup.py +20 -28
- arraybridge/oom_recovery.py +8 -8
- arraybridge/slice_processing.py +7 -3
- arraybridge/stack_utils.py +28 -31
- arraybridge/types.py +2 -25
- arraybridge/utils.py +44 -30
- {arraybridge-0.2.0.dist-info → arraybridge-0.2.8.dist-info}/METADATA +20 -5
- arraybridge-0.2.8.dist-info/RECORD +18 -0
- {arraybridge-0.2.0.dist-info → arraybridge-0.2.8.dist-info}/WHEEL +1 -1
- arraybridge/conversion_helpers.py +0 -150
- arraybridge-0.2.0.dist-info/RECORD +0 -18
- {arraybridge-0.2.0.dist-info → arraybridge-0.2.8.dist-info}/licenses/LICENSE +0 -0
arraybridge/__init__.py
CHANGED
|
@@ -5,13 +5,29 @@ This package provides automatic memory type conversion, declarative decorators,
|
|
|
5
5
|
and unified utilities for working with multiple array/tensor frameworks.
|
|
6
6
|
"""
|
|
7
7
|
|
|
8
|
-
__version__ = "0.2.
|
|
8
|
+
__version__ = "0.2.8"
|
|
9
9
|
|
|
10
10
|
from .converters import convert_memory, detect_memory_type
|
|
11
|
-
from .decorators import
|
|
11
|
+
from .decorators import (
|
|
12
|
+
DtypeConversion,
|
|
13
|
+
cupy,
|
|
14
|
+
jax,
|
|
15
|
+
memory_types,
|
|
16
|
+
numpy,
|
|
17
|
+
pyclesperanto,
|
|
18
|
+
tensorflow,
|
|
19
|
+
torch,
|
|
20
|
+
)
|
|
21
|
+
from .dtype_scaling import SCALING_FUNCTIONS
|
|
12
22
|
from .exceptions import MemoryConversionError
|
|
23
|
+
from .framework_config import _FRAMEWORK_CONFIG
|
|
24
|
+
from .framework_ops import _FRAMEWORK_OPS
|
|
25
|
+
from .gpu_cleanup import cleanup_all_gpu_frameworks
|
|
26
|
+
from .oom_recovery import _execute_with_oom_recovery
|
|
27
|
+
from .slice_processing import process_slices
|
|
13
28
|
from .stack_utils import stack_slices, unstack_slices
|
|
14
29
|
from .types import CPU_MEMORY_TYPES, GPU_MEMORY_TYPES, SUPPORTED_MEMORY_TYPES, MemoryType
|
|
30
|
+
from .utils import _ensure_module, _get_device_id, _supports_dlpack
|
|
15
31
|
|
|
16
32
|
__all__ = [
|
|
17
33
|
# Types
|
|
@@ -29,9 +45,26 @@ __all__ = [
|
|
|
29
45
|
"torch",
|
|
30
46
|
"tensorflow",
|
|
31
47
|
"jax",
|
|
48
|
+
"pyclesperanto",
|
|
49
|
+
"DtypeConversion",
|
|
32
50
|
# Stack utilities
|
|
33
51
|
"stack_slices",
|
|
34
52
|
"unstack_slices",
|
|
53
|
+
# Slice processing
|
|
54
|
+
"process_slices",
|
|
55
|
+
# GPU cleanup
|
|
56
|
+
"cleanup_all_gpu_frameworks",
|
|
35
57
|
# Exceptions
|
|
36
58
|
"MemoryConversionError",
|
|
59
|
+
# Scaling
|
|
60
|
+
"SCALING_FUNCTIONS",
|
|
61
|
+
# Framework config (internal but needed by some consumers)
|
|
62
|
+
"_FRAMEWORK_CONFIG",
|
|
63
|
+
"_FRAMEWORK_OPS",
|
|
64
|
+
# OOM recovery
|
|
65
|
+
"_execute_with_oom_recovery",
|
|
66
|
+
# Utils
|
|
67
|
+
"_ensure_module",
|
|
68
|
+
"_supports_dlpack",
|
|
69
|
+
"_get_device_id",
|
|
37
70
|
]
|
arraybridge/converters.py
CHANGED
|
@@ -4,9 +4,9 @@ from typing import Any
|
|
|
4
4
|
|
|
5
5
|
import numpy as np
|
|
6
6
|
|
|
7
|
-
from arraybridge.
|
|
7
|
+
from arraybridge.converters_registry import get_converter
|
|
8
8
|
from arraybridge.framework_config import _FRAMEWORK_CONFIG
|
|
9
|
-
from arraybridge.types import MemoryType
|
|
9
|
+
from arraybridge.types import MemoryType, VALID_MEMORY_TYPES
|
|
10
10
|
|
|
11
11
|
|
|
12
12
|
def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -> Any:
|
|
@@ -26,8 +26,14 @@ def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -
|
|
|
26
26
|
ValueError: If source_type or target_type is invalid
|
|
27
27
|
MemoryConversionError: If conversion fails
|
|
28
28
|
"""
|
|
29
|
-
|
|
30
|
-
|
|
29
|
+
if isinstance(target_type, MemoryType):
|
|
30
|
+
target_type = target_type.value
|
|
31
|
+
if target_type not in VALID_MEMORY_TYPES:
|
|
32
|
+
raise ValueError(
|
|
33
|
+
f"Invalid target_type '{target_type}'. Available types: {sorted(VALID_MEMORY_TYPES)}"
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
converter = get_converter(source_type) # Will raise ValueError if invalid
|
|
31
37
|
method = getattr(converter, f"to_{target_type}")
|
|
32
38
|
return method(data, gpu_id)
|
|
33
39
|
|
|
@@ -52,10 +58,14 @@ def detect_memory_type(data: Any) -> str:
|
|
|
52
58
|
# Check all frameworks using their module names from config
|
|
53
59
|
module_name = type(data).__module__
|
|
54
60
|
|
|
61
|
+
top_level = module_name.split(".")[0]
|
|
62
|
+
|
|
55
63
|
for mem_type, config in _FRAMEWORK_CONFIG.items():
|
|
56
|
-
import_name = config[
|
|
57
|
-
|
|
58
|
-
if
|
|
64
|
+
import_name = config["import_name"]
|
|
65
|
+
aliases = {import_name}
|
|
66
|
+
if import_name == "jax":
|
|
67
|
+
aliases.add("jaxlib")
|
|
68
|
+
if top_level in aliases:
|
|
59
69
|
return mem_type.value
|
|
60
70
|
|
|
61
71
|
raise ValueError(f"Unknown memory type for {type(data)} (module: {module_name})")
|
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Registry-based converter infrastructure using metaclass-registry.
|
|
3
|
+
|
|
4
|
+
This module provides the ConverterBase class using AutoRegisterMeta,
|
|
5
|
+
concrete converter implementations for each framework, and a helper
|
|
6
|
+
function for registry lookups.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import logging
|
|
10
|
+
from abc import abstractmethod
|
|
11
|
+
|
|
12
|
+
from metaclass_registry import AutoRegisterMeta
|
|
13
|
+
|
|
14
|
+
from arraybridge.framework_config import _FRAMEWORK_CONFIG
|
|
15
|
+
from arraybridge.types import MemoryType
|
|
16
|
+
|
|
17
|
+
logger = logging.getLogger(__name__)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class ConverterBase(metaclass=AutoRegisterMeta):
|
|
21
|
+
"""Base class for memory type converters using auto-registration.
|
|
22
|
+
|
|
23
|
+
Each concrete converter sets memory_type to register itself in the registry.
|
|
24
|
+
The registry key is the memory_type attribute (e.g., "numpy", "torch").
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
__registry_key__ = "memory_type"
|
|
28
|
+
__registry__ = {} # Simple dict - no lazy discovery needed (converters created dynamically below)
|
|
29
|
+
memory_type: str = None
|
|
30
|
+
|
|
31
|
+
@abstractmethod
|
|
32
|
+
def to_numpy(self, data, gpu_id):
|
|
33
|
+
"""Extract to NumPy (type-specific implementation)."""
|
|
34
|
+
pass
|
|
35
|
+
|
|
36
|
+
@abstractmethod
|
|
37
|
+
def from_numpy(self, data, gpu_id):
|
|
38
|
+
"""Create from NumPy (type-specific implementation)."""
|
|
39
|
+
pass
|
|
40
|
+
|
|
41
|
+
@abstractmethod
|
|
42
|
+
def from_dlpack(self, data, gpu_id):
|
|
43
|
+
"""Create from DLPack capsule (type-specific implementation)."""
|
|
44
|
+
pass
|
|
45
|
+
|
|
46
|
+
@abstractmethod
|
|
47
|
+
def move_to_device(self, data, gpu_id):
|
|
48
|
+
"""Move data to specified GPU device if needed (type-specific implementation)."""
|
|
49
|
+
pass
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _ensure_module(memory_type: str):
|
|
53
|
+
"""Import and return the module for the given memory type."""
|
|
54
|
+
from arraybridge.utils import _ensure_module as _ensure_module_impl
|
|
55
|
+
|
|
56
|
+
return _ensure_module_impl(memory_type)
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def _make_lambda_with_name(expr_str, mem_type, method_name):
|
|
60
|
+
"""Create a lambda from expression string and add proper __name__ for debugging.
|
|
61
|
+
|
|
62
|
+
Note: Uses eval() for dynamic code generation from trusted framework_config.py strings.
|
|
63
|
+
This is safe because:
|
|
64
|
+
1. Input strings come from _FRAMEWORK_CONFIG, not user input
|
|
65
|
+
2. Strings are defined at module load time by package maintainers
|
|
66
|
+
3. This pattern enables declarative framework configuration
|
|
67
|
+
"""
|
|
68
|
+
module_str = f'_ensure_module("{mem_type.value}")'
|
|
69
|
+
lambda_expr = f"lambda self, data, gpu_id: {expr_str.format(mod=module_str)}"
|
|
70
|
+
lambda_func = eval(lambda_expr)
|
|
71
|
+
lambda_func.__name__ = method_name
|
|
72
|
+
lambda_func.__qualname__ = f"{mem_type.value.capitalize()}Converter.{method_name}"
|
|
73
|
+
return lambda_func
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _make_not_implemented(mem_type_value, method_name):
|
|
77
|
+
"""Create a lambda that raises NotImplementedError with the correct signature."""
|
|
78
|
+
|
|
79
|
+
def not_impl(self, data, gpu_id):
|
|
80
|
+
raise NotImplementedError(f"DLPack not supported for {mem_type_value}")
|
|
81
|
+
|
|
82
|
+
not_impl.__name__ = method_name
|
|
83
|
+
not_impl.__qualname__ = f"{mem_type_value.capitalize()}Converter.{method_name}"
|
|
84
|
+
return not_impl
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
# Auto-generate converter classes for each memory type
|
|
88
|
+
def _create_converter_classes():
|
|
89
|
+
"""Create concrete converter classes for each memory type."""
|
|
90
|
+
converters = {}
|
|
91
|
+
|
|
92
|
+
for mem_type in MemoryType:
|
|
93
|
+
config = _FRAMEWORK_CONFIG[mem_type]
|
|
94
|
+
conversion_ops = config["conversion_ops"]
|
|
95
|
+
|
|
96
|
+
# Build class attributes
|
|
97
|
+
class_attrs = {
|
|
98
|
+
"memory_type": mem_type.value,
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
# Add conversion methods
|
|
102
|
+
for method_name, expr in conversion_ops.items():
|
|
103
|
+
if expr is None:
|
|
104
|
+
class_attrs[method_name] = _make_not_implemented(mem_type.value, method_name)
|
|
105
|
+
else:
|
|
106
|
+
class_attrs[method_name] = _make_lambda_with_name(expr, mem_type, method_name)
|
|
107
|
+
|
|
108
|
+
# Create the class
|
|
109
|
+
class_name = f"{mem_type.value.capitalize()}Converter"
|
|
110
|
+
converter_class = type(class_name, (ConverterBase,), class_attrs)
|
|
111
|
+
|
|
112
|
+
converters[mem_type] = converter_class
|
|
113
|
+
|
|
114
|
+
return converters
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
# Create all converter classes at module load time
|
|
118
|
+
_CONVERTER_CLASSES = _create_converter_classes()
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def get_converter(memory_type: str):
|
|
122
|
+
"""Get a converter instance for the given memory type.
|
|
123
|
+
|
|
124
|
+
Args:
|
|
125
|
+
memory_type: The memory type string (e.g., "numpy", "torch")
|
|
126
|
+
|
|
127
|
+
Returns:
|
|
128
|
+
A converter instance for the memory type
|
|
129
|
+
|
|
130
|
+
Raises:
|
|
131
|
+
ValueError: If memory type is not registered
|
|
132
|
+
"""
|
|
133
|
+
converter_class = ConverterBase.__registry__.get(memory_type)
|
|
134
|
+
if converter_class is None:
|
|
135
|
+
raise ValueError(
|
|
136
|
+
f"No converter registered for memory type '{memory_type}'. "
|
|
137
|
+
f"Available types: {sorted(ConverterBase.__registry__.keys())}"
|
|
138
|
+
)
|
|
139
|
+
return converter_class()
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
def _add_converter_methods():
|
|
143
|
+
"""Add to_X() methods to ConverterBase.
|
|
144
|
+
|
|
145
|
+
For each target memory type, generates a method like to_cupy(), to_torch(), etc.
|
|
146
|
+
that tries GPU-to-GPU conversion via DLPack first, then falls back to CPU roundtrip.
|
|
147
|
+
"""
|
|
148
|
+
from arraybridge.utils import _supports_dlpack
|
|
149
|
+
|
|
150
|
+
for target_type in MemoryType:
|
|
151
|
+
method_name = f"to_{target_type.value}"
|
|
152
|
+
|
|
153
|
+
def make_method(tgt):
|
|
154
|
+
def method(self, data, gpu_id):
|
|
155
|
+
# Try GPU-to-GPU first (DLPack)
|
|
156
|
+
if _supports_dlpack(data):
|
|
157
|
+
try:
|
|
158
|
+
target_converter = get_converter(tgt.value)
|
|
159
|
+
result = target_converter.from_dlpack(data, gpu_id)
|
|
160
|
+
return target_converter.move_to_device(result, gpu_id)
|
|
161
|
+
except Exception as e:
|
|
162
|
+
logger.warning(f"DLPack conversion failed: {e}. Using CPU roundtrip.")
|
|
163
|
+
|
|
164
|
+
# CPU roundtrip using polymorphism
|
|
165
|
+
numpy_data = self.to_numpy(data, gpu_id)
|
|
166
|
+
target_converter = get_converter(tgt.value)
|
|
167
|
+
return target_converter.from_numpy(numpy_data, gpu_id)
|
|
168
|
+
|
|
169
|
+
return method
|
|
170
|
+
|
|
171
|
+
setattr(ConverterBase, method_name, make_method(target_type))
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
def _validate_registry():
|
|
175
|
+
"""Validate that all memory types are registered."""
|
|
176
|
+
required_types = {mt.value for mt in MemoryType}
|
|
177
|
+
registered_types = set(ConverterBase.__registry__.keys())
|
|
178
|
+
|
|
179
|
+
if required_types != registered_types:
|
|
180
|
+
missing = required_types - registered_types
|
|
181
|
+
extra = registered_types - required_types
|
|
182
|
+
msg_parts = []
|
|
183
|
+
if missing:
|
|
184
|
+
msg_parts.append(f"Missing: {missing}")
|
|
185
|
+
if extra:
|
|
186
|
+
msg_parts.append(f"Extra: {extra}")
|
|
187
|
+
raise RuntimeError(f"Registry validation failed. {', '.join(msg_parts)}")
|
|
188
|
+
|
|
189
|
+
logger.debug(f"✅ Validated {len(registered_types)} memory type converters in registry")
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
# Add to_X() conversion methods after converter classes are created
|
|
193
|
+
_add_converter_methods()
|
|
194
|
+
|
|
195
|
+
# Run validation at module load time
|
|
196
|
+
_validate_registry()
|
arraybridge/decorators.py
CHANGED
|
@@ -30,20 +30,20 @@ from arraybridge.utils import optional_import
|
|
|
30
30
|
|
|
31
31
|
logger = logging.getLogger(__name__)
|
|
32
32
|
|
|
33
|
-
F = TypeVar(
|
|
33
|
+
F = TypeVar("F", bound=Callable[..., Any])
|
|
34
34
|
|
|
35
35
|
|
|
36
36
|
class DtypeConversion(Enum):
|
|
37
37
|
"""Data type conversion modes for all memory type functions."""
|
|
38
38
|
|
|
39
|
-
PRESERVE_INPUT = "preserve"
|
|
40
|
-
NATIVE_OUTPUT = "native"
|
|
41
|
-
UINT8 = "uint8"
|
|
42
|
-
UINT16 = "uint16"
|
|
43
|
-
INT16 = "int16"
|
|
44
|
-
INT32 = "int32"
|
|
45
|
-
FLOAT32 = "float32"
|
|
46
|
-
FLOAT64 = "float64"
|
|
39
|
+
PRESERVE_INPUT = "preserve" # Keep input dtype (default)
|
|
40
|
+
NATIVE_OUTPUT = "native" # Use framework's native output
|
|
41
|
+
UINT8 = "uint8" # Force uint8 (0-255 range)
|
|
42
|
+
UINT16 = "uint16" # Force uint16 (microscopy standard)
|
|
43
|
+
INT16 = "int16" # Force int16 (signed microscopy data)
|
|
44
|
+
INT32 = "int32" # Force int32 (large integer values)
|
|
45
|
+
FLOAT32 = "float32" # Force float32 (GPU performance)
|
|
46
|
+
FLOAT64 = "float64" # Force float64 (maximum precision)
|
|
47
47
|
|
|
48
48
|
@property
|
|
49
49
|
def numpy_dtype(self):
|
|
@@ -65,6 +65,7 @@ _gpu_frameworks_cache = {}
|
|
|
65
65
|
|
|
66
66
|
def _create_lazy_getter(framework_name: str):
|
|
67
67
|
"""Factory function that creates a lazy import getter for a framework."""
|
|
68
|
+
|
|
68
69
|
def getter():
|
|
69
70
|
if framework_name not in _gpu_frameworks_cache:
|
|
70
71
|
_gpu_frameworks_cache[framework_name] = optional_import(framework_name)
|
|
@@ -74,20 +75,22 @@ def _create_lazy_getter(framework_name: str):
|
|
|
74
75
|
f"{threading.current_thread().name}"
|
|
75
76
|
)
|
|
76
77
|
return _gpu_frameworks_cache[framework_name]
|
|
78
|
+
|
|
77
79
|
return getter
|
|
78
80
|
|
|
79
81
|
|
|
80
82
|
# Auto-generate lazy getters for all GPU frameworks
|
|
81
83
|
for mem_type in MemoryType:
|
|
82
84
|
ops = _FRAMEWORK_OPS[mem_type]
|
|
83
|
-
if ops[
|
|
84
|
-
getter_func = _create_lazy_getter(ops[
|
|
85
|
+
if ops["lazy_getter"] is not None:
|
|
86
|
+
getter_func = _create_lazy_getter(ops["import_name"])
|
|
85
87
|
globals()[f"_get_{ops['import_name']}"] = getter_func
|
|
86
88
|
|
|
87
89
|
|
|
88
90
|
# Thread-local storage for GPU streams and contexts
|
|
89
91
|
_thread_gpu_contexts = threading.local()
|
|
90
92
|
|
|
93
|
+
|
|
91
94
|
class ThreadGPUContext:
|
|
92
95
|
"""Thread-local GPU context manager for CUDA streams."""
|
|
93
96
|
|
|
@@ -100,8 +103,8 @@ class ThreadGPUContext:
|
|
|
100
103
|
def get_cupy_stream(self):
|
|
101
104
|
"""Get or create thread-local CuPy stream."""
|
|
102
105
|
if self.cupy_stream is None:
|
|
103
|
-
cupy = globals().get(
|
|
104
|
-
if cupy is not None and hasattr(cupy,
|
|
106
|
+
cupy = globals().get("_get_cupy", lambda: None)() # noqa: F821
|
|
107
|
+
if cupy is not None and hasattr(cupy, "cuda"):
|
|
105
108
|
self.cupy_stream = cupy.cuda.Stream()
|
|
106
109
|
logger.debug(f"🔧 Created CuPy stream for thread {threading.current_thread().name}")
|
|
107
110
|
return self.cupy_stream
|
|
@@ -109,33 +112,31 @@ class ThreadGPUContext:
|
|
|
109
112
|
def get_torch_stream(self):
|
|
110
113
|
"""Get or create thread-local PyTorch stream."""
|
|
111
114
|
if self.torch_stream is None:
|
|
112
|
-
torch = globals().get(
|
|
113
|
-
if torch is not None and hasattr(torch,
|
|
115
|
+
torch = globals().get("_get_torch", lambda: None)() # noqa: F821
|
|
116
|
+
if torch is not None and hasattr(torch, "cuda") and torch.cuda.is_available():
|
|
114
117
|
self.torch_stream = torch.cuda.Stream()
|
|
115
118
|
logger.debug(
|
|
116
|
-
f"🔧 Created PyTorch stream for thread "
|
|
117
|
-
f"{threading.current_thread().name}"
|
|
119
|
+
f"🔧 Created PyTorch stream for thread " f"{threading.current_thread().name}"
|
|
118
120
|
)
|
|
119
121
|
return self.torch_stream
|
|
120
122
|
|
|
121
123
|
|
|
122
124
|
def _get_thread_gpu_context():
|
|
123
125
|
"""Get or create thread-local GPU context."""
|
|
124
|
-
if not hasattr(_thread_gpu_contexts,
|
|
126
|
+
if not hasattr(_thread_gpu_contexts, "context"):
|
|
125
127
|
_thread_gpu_contexts.context = ThreadGPUContext()
|
|
126
128
|
return _thread_gpu_contexts.context
|
|
127
129
|
|
|
128
130
|
|
|
129
131
|
def memory_types(
|
|
130
|
-
input_type: str,
|
|
131
|
-
output_type: str,
|
|
132
|
-
contract: Optional[Callable[[Any], bool]] = None
|
|
132
|
+
input_type: str, output_type: str, contract: Optional[Callable[[Any], bool]] = None
|
|
133
133
|
) -> Callable[[F], F]:
|
|
134
134
|
"""
|
|
135
135
|
Base decorator for declaring memory types of a function.
|
|
136
136
|
|
|
137
137
|
This is the foundation decorator that all memory-type-specific decorators build upon.
|
|
138
138
|
"""
|
|
139
|
+
|
|
139
140
|
def decorator(func: F) -> F:
|
|
140
141
|
@functools.wraps(func)
|
|
141
142
|
def wrapper(*args, **kwargs):
|
|
@@ -171,40 +172,47 @@ def _create_dtype_wrapper(func, mem_type: MemoryType, func_name: str):
|
|
|
171
172
|
if dtype_conversion is None:
|
|
172
173
|
dtype_conversion = DtypeConversion.PRESERVE_INPUT
|
|
173
174
|
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
175
|
+
# Store original dtype
|
|
176
|
+
original_dtype = getattr(image, "dtype", None)
|
|
177
|
+
|
|
178
|
+
# Handle slice_by_slice processing for 3D arrays
|
|
179
|
+
if slice_by_slice and hasattr(image, "ndim") and image.ndim == 3:
|
|
180
|
+
result = process_slices(image, func, args, kwargs)
|
|
181
|
+
else:
|
|
182
|
+
# Call the original function normally
|
|
183
|
+
result = func(image, *args, **kwargs)
|
|
184
|
+
|
|
185
|
+
def _apply_dtype_conversion(array):
|
|
186
|
+
if dtype_conversion is None or not hasattr(array, "dtype"):
|
|
187
|
+
return array
|
|
188
|
+
if dtype_conversion == DtypeConversion.PRESERVE_INPUT:
|
|
189
|
+
# Preserve input dtype
|
|
190
|
+
if original_dtype is not None and array.dtype != original_dtype:
|
|
191
|
+
return scale_func(array, original_dtype)
|
|
192
|
+
return array
|
|
193
|
+
if dtype_conversion == DtypeConversion.NATIVE_OUTPUT:
|
|
194
|
+
# Return framework's native output dtype
|
|
195
|
+
return array
|
|
196
|
+
# Force specific dtype
|
|
197
|
+
target_dtype = dtype_conversion.numpy_dtype
|
|
198
|
+
if target_dtype is not None:
|
|
199
|
+
return scale_func(array, target_dtype)
|
|
200
|
+
return array
|
|
199
201
|
|
|
200
|
-
|
|
202
|
+
try:
|
|
203
|
+
# Apply dtype conversion to the main output
|
|
204
|
+
if isinstance(result, tuple):
|
|
205
|
+
if not result:
|
|
206
|
+
return result
|
|
207
|
+
converted_main = _apply_dtype_conversion(result[0])
|
|
208
|
+
return (converted_main, *result[1:])
|
|
209
|
+
return _apply_dtype_conversion(result)
|
|
201
210
|
except Exception as e:
|
|
202
211
|
logger.error(
|
|
203
|
-
f"Error in {mem_type.value} dtype/slice preserving wrapper "
|
|
204
|
-
f"for {func_name}: {e}"
|
|
212
|
+
f"Error in {mem_type.value} dtype/slice preserving wrapper " f"for {func_name}: {e}"
|
|
205
213
|
)
|
|
206
|
-
# Return
|
|
207
|
-
return
|
|
214
|
+
# Return unmodified result on conversion errors
|
|
215
|
+
return result
|
|
208
216
|
|
|
209
217
|
# Update function signature to include new parameters
|
|
210
218
|
try:
|
|
@@ -215,22 +223,19 @@ def _create_dtype_wrapper(func, mem_type: MemoryType, func_name: str):
|
|
|
215
223
|
param_names = [p.name for p in new_params]
|
|
216
224
|
|
|
217
225
|
# Add dtype_conversion parameter first (before slice_by_slice)
|
|
218
|
-
if
|
|
226
|
+
if "dtype_conversion" not in param_names:
|
|
219
227
|
dtype_param = inspect.Parameter(
|
|
220
|
-
|
|
228
|
+
"dtype_conversion",
|
|
221
229
|
inspect.Parameter.KEYWORD_ONLY,
|
|
222
230
|
default=DtypeConversion.PRESERVE_INPUT,
|
|
223
|
-
annotation=Optional[DtypeConversion]
|
|
231
|
+
annotation=Optional[DtypeConversion],
|
|
224
232
|
)
|
|
225
233
|
new_params.append(dtype_param)
|
|
226
234
|
|
|
227
235
|
# Add slice_by_slice parameter
|
|
228
|
-
if
|
|
236
|
+
if "slice_by_slice" not in param_names:
|
|
229
237
|
slice_param = inspect.Parameter(
|
|
230
|
-
|
|
231
|
-
inspect.Parameter.KEYWORD_ONLY,
|
|
232
|
-
default=False,
|
|
233
|
-
annotation=bool
|
|
238
|
+
"slice_by_slice", inspect.Parameter.KEYWORD_ONLY, default=False, annotation=bool
|
|
234
239
|
)
|
|
235
240
|
new_params.append(slice_param)
|
|
236
241
|
|
|
@@ -241,23 +246,18 @@ def _create_dtype_wrapper(func, mem_type: MemoryType, func_name: str):
|
|
|
241
246
|
# Update docstring
|
|
242
247
|
if dtype_wrapper.__doc__:
|
|
243
248
|
dtype_wrapper.__doc__ += (
|
|
244
|
-
f"\n\n Additional Parameters "
|
|
245
|
-
f"(added by {mem_type.value} decorator):\n"
|
|
249
|
+
f"\n\n Additional Parameters " f"(added by {mem_type.value} decorator):\n"
|
|
246
250
|
)
|
|
247
251
|
dtype_wrapper.__doc__ += (
|
|
248
252
|
" dtype_conversion (DtypeConversion, optional): "
|
|
249
253
|
"How to handle output dtype.\n"
|
|
250
254
|
)
|
|
255
|
+
dtype_wrapper.__doc__ += " Defaults to PRESERVE_INPUT (match input dtype).\n"
|
|
251
256
|
dtype_wrapper.__doc__ += (
|
|
252
|
-
"
|
|
257
|
+
" slice_by_slice (bool, optional): " "Process 3D arrays slice-by-slice.\n"
|
|
253
258
|
)
|
|
254
259
|
dtype_wrapper.__doc__ += (
|
|
255
|
-
"
|
|
256
|
-
"Process 3D arrays slice-by-slice.\n"
|
|
257
|
-
)
|
|
258
|
-
dtype_wrapper.__doc__ += (
|
|
259
|
-
" Defaults to False. "
|
|
260
|
-
"Prevents cross-slice contamination.\n"
|
|
260
|
+
" Defaults to False. " "Prevents cross-slice contamination.\n"
|
|
261
261
|
)
|
|
262
262
|
|
|
263
263
|
except Exception as e:
|
|
@@ -273,8 +273,8 @@ def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool):
|
|
|
273
273
|
This function creates the GPU-specific wrapper with stream management and OOM recovery.
|
|
274
274
|
"""
|
|
275
275
|
ops = _FRAMEWORK_OPS[mem_type]
|
|
276
|
-
framework_name = ops[
|
|
277
|
-
lazy_getter = globals().get(ops[
|
|
276
|
+
framework_name = ops["import_name"]
|
|
277
|
+
lazy_getter = globals().get(ops["lazy_getter"])
|
|
278
278
|
|
|
279
279
|
@functools.wraps(func)
|
|
280
280
|
def gpu_wrapper(*args, **kwargs):
|
|
@@ -282,7 +282,7 @@ def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool):
|
|
|
282
282
|
|
|
283
283
|
# Check if GPU is available for this framework
|
|
284
284
|
if framework is not None:
|
|
285
|
-
gpu_check_expr = ops[
|
|
285
|
+
gpu_check_expr = ops["gpu_check"].format(mod=framework_name)
|
|
286
286
|
try:
|
|
287
287
|
gpu_available = eval(gpu_check_expr, {framework_name: framework})
|
|
288
288
|
except Exception:
|
|
@@ -308,7 +308,7 @@ def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool):
|
|
|
308
308
|
return func(*args, **kwargs)
|
|
309
309
|
|
|
310
310
|
# Execute with OOM recovery if enabled
|
|
311
|
-
if oom_recovery and ops[
|
|
311
|
+
if oom_recovery and ops["has_oom_recovery"]:
|
|
312
312
|
return _execute_with_oom_recovery(execute_with_stream, mem_type.value)
|
|
313
313
|
else:
|
|
314
314
|
return execute_with_stream()
|
|
@@ -331,8 +331,14 @@ def _create_memory_decorator(mem_type: MemoryType):
|
|
|
331
331
|
"""
|
|
332
332
|
ops = _FRAMEWORK_OPS[mem_type]
|
|
333
333
|
|
|
334
|
-
def decorator(
|
|
335
|
-
|
|
334
|
+
def decorator(
|
|
335
|
+
func=None,
|
|
336
|
+
*,
|
|
337
|
+
input_type=mem_type.value,
|
|
338
|
+
output_type=mem_type.value,
|
|
339
|
+
oom_recovery=True,
|
|
340
|
+
contract=None,
|
|
341
|
+
):
|
|
336
342
|
"""
|
|
337
343
|
Decorator for {mem_type} memory type functions.
|
|
338
344
|
|
|
@@ -346,12 +352,11 @@ def _create_memory_decorator(mem_type: MemoryType):
|
|
|
346
352
|
Returns:
|
|
347
353
|
Decorated function with memory type metadata and dtype preservation
|
|
348
354
|
"""
|
|
355
|
+
|
|
349
356
|
def inner_decorator(func):
|
|
350
357
|
# Apply base memory_types decorator
|
|
351
358
|
memory_decorator = memory_types(
|
|
352
|
-
input_type=input_type,
|
|
353
|
-
output_type=output_type,
|
|
354
|
-
contract=contract
|
|
359
|
+
input_type=input_type, output_type=output_type, contract=contract
|
|
355
360
|
)
|
|
356
361
|
func = memory_decorator(func)
|
|
357
362
|
|
|
@@ -359,7 +364,7 @@ def _create_memory_decorator(mem_type: MemoryType):
|
|
|
359
364
|
func = _create_dtype_wrapper(func, mem_type, func.__name__)
|
|
360
365
|
|
|
361
366
|
# Apply GPU wrapper if this is a GPU memory type
|
|
362
|
-
if ops[
|
|
367
|
+
if ops["gpu_check"] is not None:
|
|
363
368
|
func = _create_gpu_wrapper(func, mem_type, oom_recovery)
|
|
364
369
|
|
|
365
370
|
return func
|
|
@@ -371,7 +376,7 @@ def _create_memory_decorator(mem_type: MemoryType):
|
|
|
371
376
|
|
|
372
377
|
# Set proper function name and docstring
|
|
373
378
|
decorator.__name__ = mem_type.value
|
|
374
|
-
decorator.__doc__ = decorator.__doc__.format(mem_type=ops[
|
|
379
|
+
decorator.__doc__ = decorator.__doc__.format(mem_type=ops["display_name"])
|
|
375
380
|
|
|
376
381
|
return decorator
|
|
377
382
|
|
|
@@ -384,13 +389,12 @@ for mem_type in MemoryType:
|
|
|
384
389
|
|
|
385
390
|
# Export all decorators
|
|
386
391
|
__all__ = [
|
|
387
|
-
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
392
|
+
"memory_types",
|
|
393
|
+
"DtypeConversion",
|
|
394
|
+
"numpy", # noqa: F822
|
|
395
|
+
"cupy", # noqa: F822
|
|
396
|
+
"torch", # noqa: F822
|
|
397
|
+
"tensorflow", # noqa: F822
|
|
398
|
+
"jax", # noqa: F822
|
|
399
|
+
"pyclesperanto", # noqa: F822
|
|
395
400
|
]
|
|
396
|
-
|