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 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.0"
8
+ __version__ = "0.2.8"
9
9
 
10
10
  from .converters import convert_memory, detect_memory_type
11
- from .decorators import cupy, jax, memory_types, numpy, tensorflow, torch
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.conversion_helpers import _CONVERTERS
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
- source_enum = MemoryType(source_type)
30
- converter = _CONVERTERS[source_enum]
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['import_name']
57
- # Check if module name starts with or contains the import name
58
- if module_name.startswith(import_name) or import_name in module_name:
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('F', bound=Callable[..., Any])
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" # 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)
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['lazy_getter'] is not None:
84
- getter_func = _create_lazy_getter(ops['import_name'])
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('_get_cupy', lambda: None)() # noqa: F821
104
- if cupy is not None and hasattr(cupy, 'cuda'):
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('_get_torch', lambda: None)() # noqa: F821
113
- if torch is not None and hasattr(torch, 'cuda') and torch.cuda.is_available():
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, 'context'):
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
- try:
175
- # Store original dtype
176
- original_dtype = image.dtype
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
- # Apply dtype conversion based on enum value
186
- if hasattr(result, 'dtype') and dtype_conversion is not None:
187
- if dtype_conversion == DtypeConversion.PRESERVE_INPUT:
188
- # Preserve input dtype
189
- if result.dtype != original_dtype:
190
- result = scale_func(result, original_dtype)
191
- elif dtype_conversion == DtypeConversion.NATIVE_OUTPUT:
192
- # Return framework's native output dtype
193
- pass # No conversion needed
194
- else:
195
- # Force specific dtype
196
- target_dtype = dtype_conversion.numpy_dtype
197
- if target_dtype is not None:
198
- result = scale_func(result, target_dtype)
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
- return result
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 original result on error
207
- return func(image, *args, **kwargs)
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 'dtype_conversion' not in param_names:
226
+ if "dtype_conversion" not in param_names:
219
227
  dtype_param = inspect.Parameter(
220
- 'dtype_conversion',
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 'slice_by_slice' not in param_names:
236
+ if "slice_by_slice" not in param_names:
229
237
  slice_param = inspect.Parameter(
230
- 'slice_by_slice',
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
- " Defaults to PRESERVE_INPUT (match input dtype).\n"
257
+ " slice_by_slice (bool, optional): " "Process 3D arrays slice-by-slice.\n"
253
258
  )
254
259
  dtype_wrapper.__doc__ += (
255
- " slice_by_slice (bool, optional): "
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['import_name']
277
- lazy_getter = globals().get(ops['lazy_getter'])
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['gpu_check'].format(mod=framework_name)
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['has_oom_recovery']:
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(func=None, *, input_type=mem_type.value, output_type=mem_type.value,
335
- oom_recovery=True, contract=None):
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['gpu_check'] is not None:
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['display_name'])
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
- 'memory_types',
388
- 'DtypeConversion',
389
- 'numpy', # noqa: F822
390
- 'cupy', # noqa: F822
391
- 'torch', # noqa: F822
392
- 'tensorflow', # noqa: F822
393
- 'jax', # noqa: F822
394
- 'pyclesperanto', # noqa: F822
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
-