Coverage for src/arraybridge/__init__.py: 100%

22 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 11:15 +0000

1""" 

2arraybridge: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto. 

3 

4This package provides automatic memory type conversion, declarative decorators, 

5and unified utilities for working with multiple array/tensor frameworks. 

6""" 

7 

8__version__ = "0.3.3" 

9 

10from . import decorators as _decorators 

11from .array_geometry import ArrayGeometry 

12from .array_payload import ArrayPayload 

13from .converters import convert_memory, detect_memory_type 

14from .dtype_scaling import SCALING_FUNCTIONS 

15from .exceptions import MemoryConversionError 

16from .framework_config import _FRAMEWORK_CONFIG 

17from .framework_ops import _FRAMEWORK_OPS 

18from .gpu_cleanup import cleanup_all_gpu_frameworks 

19from .oom_recovery import _execute_with_oom_recovery 

20from .slice_processing import process_slices 

21from .stack_utils import stack_slices, unstack_slices 

22from .types import ( 

23 CPU_MEMORY_TYPES, 

24 GPU_MEMORY_TYPES, 

25 SUPPORTED_MEMORY_TYPES, 

26 MemoryContractAttribute, 

27 MemoryType, 

28) 

29from .utils import _ensure_module, _get_device_id, _supports_dlpack 

30 

31DtypeConversion = _decorators.DtypeConversion 

32SliceBySliceRuntimeParameter = _decorators.SliceBySliceRuntimeParameter 

33memory_types = _decorators.memory_types 

34wrap_dtype_preserving_callable = _decorators.wrap_dtype_preserving_callable 

35for _memory_type in MemoryType: 

36 globals()[_memory_type.value] = getattr(_decorators, _memory_type.value) 

37 

38__all__ = [ 

39 # Types 

40 "MemoryType", 

41 "MemoryContractAttribute", 

42 "ArrayPayload", 

43 "ArrayGeometry", 

44 "CPU_MEMORY_TYPES", 

45 "GPU_MEMORY_TYPES", 

46 "SUPPORTED_MEMORY_TYPES", 

47 # Converters 

48 "convert_memory", 

49 "detect_memory_type", 

50 # Decorators 

51 "memory_types", 

52 "DtypeConversion", 

53 "SliceBySliceRuntimeParameter", 

54 "wrap_dtype_preserving_callable", 

55 # Stack utilities 

56 "stack_slices", 

57 "unstack_slices", 

58 # Slice processing 

59 "process_slices", 

60 # GPU cleanup 

61 "cleanup_all_gpu_frameworks", 

62 # Exceptions 

63 "MemoryConversionError", 

64 # Scaling 

65 "SCALING_FUNCTIONS", 

66 # Framework config (internal but needed by some consumers) 

67 "_FRAMEWORK_CONFIG", 

68 "_FRAMEWORK_OPS", 

69 # OOM recovery 

70 "_execute_with_oom_recovery", 

71 # Utils 

72 "_ensure_module", 

73 "_supports_dlpack", 

74 "_get_device_id", 

75] + [memory_type.value for memory_type in MemoryType]