arraybridge 0.3.3__py3-none-any.whl → 0.3.6__py3-none-any.whl

arraybridge/__init__.py

@@ -3,37 +3,49 @@ arraybridge: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclespe
 
 This package provides automatic memory type conversion, declarative decorators,
 and unified utilities for working with multiple array/tensor frameworks.
-"""
 
-__version__ = "0.3.3"
+Package-root imports stay lightweight: submodule exports load on first
+attribute access instead of at import time. Declaration-only consumers
+(memory-type tables, configuration modules) therefore do not pay the
+NumPy/numcodecs import cost.
+"""
 
-from . import decorators as _decorators
-from .array_geometry import ArrayGeometry
-from .array_payload import ArrayPayload
-from .converters import convert_memory, detect_memory_type
-from .dtype_scaling import SCALING_FUNCTIONS
-from .exceptions import MemoryConversionError
-from .framework_config import _FRAMEWORK_CONFIG
-from .framework_ops import _FRAMEWORK_OPS
-from .gpu_cleanup import cleanup_all_gpu_frameworks
-from .oom_recovery import _execute_with_oom_recovery
-from .slice_processing import process_slices
-from .stack_utils import stack_slices, unstack_slices
-from .types import (
-    CPU_MEMORY_TYPES,
-    GPU_MEMORY_TYPES,
-    SUPPORTED_MEMORY_TYPES,
-    MemoryContractAttribute,
-    MemoryType,
-)
-from .utils import _ensure_module, _get_device_id, _supports_dlpack
+__version__ = "0.3.6"
 
-DtypeConversion = _decorators.DtypeConversion
-SliceBySliceRuntimeParameter = _decorators.SliceBySliceRuntimeParameter
-memory_types = _decorators.memory_types
-wrap_dtype_preserving_callable = _decorators.wrap_dtype_preserving_callable
-for _memory_type in MemoryType:
-    globals()[_memory_type.value] = getattr(_decorators, _memory_type.value)
+_LAZY_EXPORTS: dict[str, str] = {
+    "MemoryType": ".types",
+    "MemoryContractAttribute": ".types",
+    "ArrayPayload": ".array_payload",
+    "ArrayGeometry": ".array_geometry",
+    "CPU_MEMORY_TYPES": ".types",
+    "GPU_MEMORY_TYPES": ".types",
+    "SUPPORTED_MEMORY_TYPES": ".types",
+    "convert_memory": ".converters",
+    "detect_memory_type": ".converters",
+    "memory_types": ".decorators",
+    "DtypeConversion": ".decorators",
+    "SliceBySliceRuntimeParameter": ".decorators",
+    "wrap_dtype_preserving_callable": ".decorators",
+    "stack_slices": ".stack_utils",
+    "unstack_slices": ".stack_utils",
+    "process_slices": ".slice_processing",
+    "cleanup_all_gpu_frameworks": ".gpu_cleanup",
+    "MemoryConversionError": ".exceptions",
+    "SCALING_FUNCTIONS": ".dtype_scaling",
+    "_FRAMEWORK_CONFIG": ".framework_config",
+    "_FRAMEWORK_OPS": ".framework_ops",
+    "_execute_with_oom_recovery": ".oom_recovery",
+    "_ensure_module": ".utils",
+    "_supports_dlpack": ".utils",
+    "_get_device_id": ".utils",
+    # Decorator exports named after each memory type (numpy, cupy, torch, ...).
+    "numpy": ".decorators",
+    "cupy": ".decorators",
+    "torch": ".decorators",
+    "tensorflow": ".decorators",
+    "jax": ".decorators",
+    "pyclesperanto": ".decorators",
+}
 
 __all__ = [
     # Types
@@ -72,4 +84,26 @@ __all__ = [
     "_ensure_module",
     "_supports_dlpack",
     "_get_device_id",
-] + [memory_type.value for memory_type in MemoryType]
+    "numpy",
+    "cupy",
+    "torch",
+    "tensorflow",
+    "jax",
+    "pyclesperanto",
+]
+
+
+def __getattr__(name: str):
+    module_name = _LAZY_EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    import importlib
+
+    value = getattr(importlib.import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY_EXPORTS))

arraybridge/decorators.py

@@ -82,11 +82,13 @@ class DtypeConversionConfig(RuntimeParameterDeclarationABC):
         return DtypeConversionConfig
 
     @classmethod
-    def parameter(cls) -> inspect.Parameter:
+    def parameter(
+        cls, *, default_value: "DtypeConversionConfig | None" = None
+    ) -> inspect.Parameter:
         return inspect.Parameter(
             cls.require_parameter_name(),
             inspect.Parameter.KEYWORD_ONLY,
-            default=cls.default_value(),
+            default=cls.default_value() if default_value is None else default_value,
             annotation=cls.annotation_type(),
         )
 
@@ -264,12 +266,30 @@ class KeywordOnlySignatureExtension:
         return len(parameters)
 
 
-# Thread-local storage for GPU streams and contexts
-_thread_gpu_contexts = threading.local()
+class ThreadGPUContextStorage(threading.local):
+    """Declare the existing native thread-local context slot and its type."""
+
+    context: "ThreadGPUContext | None" = None
 
 
 class ThreadGPUContext:
-    """Thread-local streams keyed by framework-local device identity."""
+    """Runtime-owned thread-local streams keyed by framework/device identity.
+
+    Keep the runtime handle on its importable owner, not in decorator function
+    globals. A retained, unpublished decorated callable is serialized by value;
+    its durable closure must not pull a ``threading.local`` into history.
+    """
+
+    _contexts: ClassVar[ThreadGPUContextStorage] = ThreadGPUContextStorage()
+
+    @classmethod
+    def current(cls) -> "ThreadGPUContext":
+        """Return this thread's runtime context without serializing its handle."""
+        context = cls._contexts.context
+        if context is None:
+            context = cls()
+            cls._contexts.context = context
+        return context
 
     def __init__(self):
         self._streams: dict[tuple[MemoryType, int], Any] = {}
@@ -301,13 +321,6 @@ class ThreadGPUContext:
         return device_id, self._streams[key]
 
 
-def _get_thread_gpu_context():
-    """Get or create thread-local GPU context."""
-    if not hasattr(_thread_gpu_contexts, "context"):
-        _thread_gpu_contexts.context = ThreadGPUContext()
-    return _thread_gpu_contexts.context
-
-
 def memory_types(
     input_type: str | MemoryType,
     output_type: str | MemoryType,
@@ -352,6 +365,7 @@ def wrap_dtype_preserving_callable(
     mem_type: MemoryType,
     *,
     slice_by_slice_default: bool = False,
+    dtype_config_default: DtypeConversionConfig | None = None,
 ):
     """
     Return a callable with ArrayBridge dtype and slice controls.
@@ -363,18 +377,25 @@ def wrap_dtype_preserving_callable(
     input_memory_type = MemoryType(MemoryContractAttribute.INPUT.read(func, mem_type.value))
     output_memory_type = MemoryType(MemoryContractAttribute.OUTPUT.read(func, mem_type.value))
     scale_func = output_memory_type.scale_dtype
+    default_dtype_config = (
+        DtypeConversionConfig.default_value()
+        if dtype_config_default is None
+        else dtype_config_default
+    )
+    if not isinstance(default_dtype_config, DtypeConversionConfig):
+        raise TypeError("Callable dtype default must be a DtypeConversionConfig.")
 
     @functools.wraps(func)
     def dtype_wrapper(image, *args, **kwargs):
         # Pipeline runtimes may inject dtype_config; direct calls use the same
-        # preserve-input default explicitly.
+        # callable-owned default explicitly (preserve-input when undeclared).
         slice_by_slice = kwargs.pop(
             SliceBySliceRuntimeParameter.require_parameter_name(),
             slice_by_slice_default,
         )
         dtype_config: DtypeConversionConfig = kwargs.pop(
             DtypeConversionConfig.require_parameter_name(),
-            DtypeConversionConfig.default_value(),
+            default_dtype_config,
         )
         dtype_conversion = dtype_config.default_dtype_conversion
 
@@ -424,7 +445,7 @@ def wrap_dtype_preserving_callable(
             )
         )
         dtype_signature = KeywordOnlySignatureExtension(dtype_signature).with_parameter(
-            DtypeConversionConfig.parameter()
+            DtypeConversionConfig.parameter(default_value=default_dtype_config)
         )
         setattr(dtype_wrapper, "__signature__", dtype_signature)
 
@@ -461,7 +482,7 @@ def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool):
         # Check if GPU is available for this framework
         if framework is not None and mem_type.available_device_ids(framework):
             # Get thread-local context
-            ctx = _get_thread_gpu_context()
+            ctx = ThreadGPUContext.current()
 
             device_id, stream = ctx.stream_for(mem_type, framework)
 
@@ -510,6 +531,7 @@ def _create_memory_decorator(mem_type: MemoryType):
         oom_recovery=True,
         contract=None,
         slice_by_slice_default=False,
+        dtype_config_default: DtypeConversionConfig | None = None,
     ):
         """
         Decorator for {mem_type} memory type functions.
@@ -521,6 +543,8 @@ def _create_memory_decorator(mem_type: MemoryType):
             oom_recovery: Enable automatic OOM recovery (default: True)
             contract: Optional validation function for outputs
             slice_by_slice_default: Default for the decorator-owned slice control
+            dtype_config_default: Callable-owned dtype policy for direct calls;
+                explicit runtime dtype_config still overrides this default.
 
         Returns:
             Decorated function with memory type metadata and dtype preservation
@@ -538,6 +562,7 @@ def _create_memory_decorator(mem_type: MemoryType):
                 func,
                 mem_type,
                 slice_by_slice_default=slice_by_slice_default,
+                dtype_config_default=dtype_config_default,
             )
 
             # Apply GPU wrapper if this is a GPU memory type

arraybridge/types.py

@@ -5,6 +5,8 @@ This module defines the MemoryType enum and related constants for managing
 different array/tensor frameworks.
 """
 
+from __future__ import annotations
+
 import importlib
 import importlib.util
 import logging
@@ -15,17 +17,10 @@ from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, TypeVar, cast
-
-from arraybridge.array_operations import (
-    CUPY_OPERATIONS,
-    JAX_OPERATIONS,
-    NUMPY_OPERATIONS,
-    PYCLESPERANTO_OPERATIONS,
-    TENSORFLOW_OPERATIONS,
-    TORCH_OPERATIONS,
-    ArrayOperations,
-)
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+if TYPE_CHECKING:
+    from arraybridge.array_operations import ArrayOperations
 
 T = TypeVar("T")
 logger = logging.getLogger(__name__)
@@ -561,7 +556,7 @@ class _MemoryTypeFields:
     module_aliases: tuple[str, ...]
     import_environment: tuple[tuple[str, str], ...]
     _runtime: FrameworkRuntime
-    _operations: ArrayOperations
+    _operations_name: str
 
 
 class MemoryType(_MemoryTypeFields, Enum):
@@ -571,7 +566,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         cls,
         value: str,
         *declaration: Any,
-    ) -> "MemoryType":
+    ) -> MemoryType:
         (
             import_name,
             display_name,
@@ -589,9 +584,18 @@ class MemoryType(_MemoryTypeFields, Enum):
         member.module_aliases = cast(tuple[str, ...], module_aliases)
         member.import_environment = cast(tuple[tuple[str, str], ...], import_environment)
         member._runtime = cast(FrameworkRuntime, runtime)
-        member._operations = cast(ArrayOperations, operations)
+        member._operations_name = cast(str, operations)
         return member
 
+    @property
+    def _operations(self) -> ArrayOperations:
+        """Resolve this declaration's operations registry on first access."""
+
+        from arraybridge import array_operations
+
+        operations: ArrayOperations = getattr(array_operations, self._operations_name)
+        return operations
+
     NUMPY = (
         "numpy",
         "numpy",
@@ -600,7 +604,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         (),
         (),
         FrameworkRuntime(oom_matcher=_numpy_oom),
-        NUMPY_OPERATIONS,
+        "NUMPY_OPERATIONS",
     )
     CUPY = (
         "cupy",
@@ -624,7 +628,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             oom_matcher=_cupy_oom,
             subprocess_environment=_nvidia_wheel_subprocess_environment,
         ),
-        CUPY_OPERATIONS,
+        "CUPY_OPERATIONS",
     )
     TORCH = (
         "torch",
@@ -647,7 +651,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_exporter=_protocol_dlpack_export,
             oom_matcher=_torch_oom,
         ),
-        TORCH_OPERATIONS,
+        "TORCH_OPERATIONS",
     )
     TENSORFLOW = (
         "tensorflow",
@@ -666,7 +670,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_validator=_tensorflow_dlpack,
             oom_matcher=_tensorflow_oom,
         ),
-        TENSORFLOW_OPERATIONS,
+        "TENSORFLOW_OPERATIONS",
     )
     JAX = (
         "jax",
@@ -684,7 +688,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             dlpack_exporter=_protocol_dlpack_export,
             oom_matcher=_jax_oom,
         ),
-        JAX_OPERATIONS,
+        "JAX_OPERATIONS",
     )
     PYCLESPERANTO = (
         "pyclesperanto",
@@ -701,7 +705,7 @@ class MemoryType(_MemoryTypeFields, Enum):
             move_to_active_device=_pyclesperanto_device_move,
             oom_matcher=_pyclesperanto_oom,
         ),
-        PYCLESPERANTO_OPERATIONS,
+        "PYCLESPERANTO_OPERATIONS",
     )
 
     @property
@@ -953,7 +957,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data)
         return importer(payload, framework)
 
-    def convert_to(self, data: Any, target: "MemoryType", device_id: int) -> Any:
+    def convert_to(self, data: Any, target: MemoryType, device_id: int) -> Any:
         """Convert one array through the source and target declarations."""
 
         if self is target:

{arraybridge-0.3.3.dist-info → arraybridge-0.3.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraybridge
-Version: 0.3.3
+Version: 0.3.6
 Summary: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion
 Project-URL: Homepage, https://github.com/OpenHCSDev/arraybridge
 Project-URL: Documentation, https://arraybridge.readthedocs.io
@@ -36,6 +36,7 @@ Requires-Dist: cupy>=10.0; extra == 'cupy'
 Provides-Extra: dev
 Requires-Dist: black>=23.0; extra == 'dev'
 Requires-Dist: build>=1.0; extra == 'dev'
+Requires-Dist: dill>=0.3.8; extra == 'dev'
 Requires-Dist: mypy>=1.0; extra == 'dev'
 Requires-Dist: packaging>=23.0; extra == 'dev'
 Requires-Dist: pytest-cov>=4.0; extra == 'dev'

{arraybridge-0.3.3.dist-info → arraybridge-0.3.6.dist-info}/RECORD

@@ -1,10 +1,10 @@
-arraybridge/__init__.py,sha256=TmDOlpJsQxwY_k9aQ6YlGav2mZLLPzmh2tWbAOY1FyQ,2326
+arraybridge/__init__.py,sha256=rhst2HuRsZE1bp4f23SjK6cTHNpjQUAsSnqgh9R6QUA,3184
 arraybridge/array_geometry.py,sha256=CopfGmFVKwV5l-lOTImhm1YVh18weUng0ReEF58qlPI,1689
 arraybridge/array_operations.py,sha256=fp69o9BYcCzfhmLcWrH68SVGSfPL3nb6Cty-Amd_oQM,13082
 arraybridge/array_payload.py,sha256=6_jJwPRya0MHrUNSA9kXPlTISpLa0G74rqDValGXKx4,1079
 arraybridge/converters.py,sha256=q44E6Pg1rJz8yq6WQIPLOZPv1lqqcUcEG4Gh4D5M1_o,2649
 arraybridge/converters_registry.py,sha256=_hHtbd-3SMIKUw8fvNu3twPWJY7qHyNyQ_NhVp3eSHQ,5647
-arraybridge/decorators.py,sha256=mYIxBbCuWQVN8dhAE2M7i2Yo0KbcMyY4MZmsfGXGm-0,19975
+arraybridge/decorators.py,sha256=KPm_R3C9TMHvNARf3vKAZQY-3gKSlvYrimxBQMyI5A4,21271
 arraybridge/dtype_scaling.py,sha256=KDgYyCvxgsQXLJJl_bs5nH46zQimsA7c1Olv6xJ2z1o,831
 arraybridge/exceptions.py,sha256=fhej5ZS39QcBFWzNeNYjoqeLVnOgwwp-cWIme25jyuo,752
 arraybridge/framework_config.py,sha256=5yTTt9NWcwzKDMXbc-fqCdiFl-piKOt3SUnxmRXXCQk,744
@@ -13,9 +13,9 @@ arraybridge/gpu_cleanup.py,sha256=ADI4wMLk9GIO9rJ6BJuNk3uZIyANoAqJ7VFTnLJzrbE,25
 arraybridge/oom_recovery.py,sha256=hGvOuJHVSopiIqblzoMFlqW8_oK6ZlAGF_l4pfEljeI,2403
 arraybridge/slice_processing.py,sha256=5R10OmHW8BSbJFwT1URd2Sz3P15o5EhaI2VnVUUIvkQ,4450
 arraybridge/stack_utils.py,sha256=FVikN0XG-Xryw-1p9MaSRUTZet7kLcy0gTEERSjc9r4,6915
-arraybridge/types.py,sha256=F_aJLJnZ46KBmCziHjtfkkvpqzHkc-kDjbsuRc98uak,35588
+arraybridge/types.py,sha256=e5R2DVflQAjXrfTk033KRdiiwoip3IQAvgDCYUjPdiQ,35807
 arraybridge/utils.py,sha256=RqThtYScPEOEQG491dA8b_VvUTjgcFhwLbh0erQu5Ns,4833
-arraybridge-0.3.3.dist-info/METADATA,sha256=eloa3YEUoJnRo0S3n6JgClhYWy8J8sdAtINIwsroOf8,5275
-arraybridge-0.3.3.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
-arraybridge-0.3.3.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
-arraybridge-0.3.3.dist-info/RECORD,,
+arraybridge-0.3.6.dist-info/METADATA,sha256=uOSPAN8b0foxSeEJslCFyv0VhJR7suC1TWL9RBCrLTo,5318
+arraybridge-0.3.6.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+arraybridge-0.3.6.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
+arraybridge-0.3.6.dist-info/RECORD,,