Decorators Module

Memory type declaration decorators for OpenHCS.

This module provides decorators for explicitly declaring the memory interface of pure functions, enforcing Clause 106-A (Declared Memory Types) and supporting memory-type-aware dispatching and orchestration.

These decorators annotate functions with input_memory_type and output_memory_type attributes and provide automatic thread-local CUDA stream management for GPU frameworks to enable true parallelization across multiple threads.

REFACTORED: Uses enum-driven metaprogramming to eliminate 79% of code duplication.

memory_types(input_type, output_type, contract=None)[source]

Base decorator for declaring memory types of a function.

This is the foundation decorator that all memory-type-specific decorators build upon.

Return type:

Callable[[TypeVar(F, bound= Callable[..., Any])], TypeVar(F, bound= Callable[..., Any])]

class DtypeConversion(value)[source]

Bases: Enum

Data type conversion modes for all memory type functions.

PRESERVE_INPUT = 'preserve'
NATIVE_OUTPUT = 'native'
UINT8 = 'uint8'
UINT16 = 'uint16'
INT16 = 'int16'
INT32 = 'int32'
FLOAT32 = 'float32'
FLOAT64 = 'float64'
property numpy_dtype

Get the corresponding numpy dtype.

numpy(func=None, *, input_type='numpy', output_type='numpy', oom_recovery=True, contract=None)

Decorator for NumPy memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: NumPy)

  • output_type – Expected output memory type (default: NumPy)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

cupy(func=None, *, input_type='cupy', output_type='cupy', oom_recovery=True, contract=None)

Decorator for CuPy memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: CuPy)

  • output_type – Expected output memory type (default: CuPy)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

torch(func=None, *, input_type='torch', output_type='torch', oom_recovery=True, contract=None)

Decorator for PyTorch memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: PyTorch)

  • output_type – Expected output memory type (default: PyTorch)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

tensorflow(func=None, *, input_type='tensorflow', output_type='tensorflow', oom_recovery=True, contract=None)

Decorator for TensorFlow memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: TensorFlow)

  • output_type – Expected output memory type (default: TensorFlow)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

jax(func=None, *, input_type='jax', output_type='jax', oom_recovery=True, contract=None)

Decorator for JAX memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: JAX)

  • output_type – Expected output memory type (default: JAX)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

pyclesperanto(func=None, *, input_type='pyclesperanto', output_type='pyclesperanto', oom_recovery=True, contract=None)

Decorator for pyclesperanto memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: pyclesperanto)

  • output_type – Expected output memory type (default: pyclesperanto)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

Available Decorators

Base Decorator

memory_types(input_type, output_type, contract=None)[source]

Base decorator for declaring memory types of a function.

This is the foundation decorator that all memory-type-specific decorators build upon.

Return type:

Callable[[TypeVar(F, bound= Callable[..., Any])], TypeVar(F, bound= Callable[..., Any])]

Framework-Specific Decorators

numpy(func=None, *, input_type='numpy', output_type='numpy', oom_recovery=True, contract=None)

Decorator for NumPy memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: NumPy)

  • output_type – Expected output memory type (default: NumPy)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

cupy(func=None, *, input_type='cupy', output_type='cupy', oom_recovery=True, contract=None)

Decorator for CuPy memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: CuPy)

  • output_type – Expected output memory type (default: CuPy)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

torch(func=None, *, input_type='torch', output_type='torch', oom_recovery=True, contract=None)

Decorator for PyTorch memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: PyTorch)

  • output_type – Expected output memory type (default: PyTorch)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

tensorflow(func=None, *, input_type='tensorflow', output_type='tensorflow', oom_recovery=True, contract=None)

Decorator for TensorFlow memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: TensorFlow)

  • output_type – Expected output memory type (default: TensorFlow)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

jax(func=None, *, input_type='jax', output_type='jax', oom_recovery=True, contract=None)

Decorator for JAX memory type functions.

Parameters:

  • func – Function to decorate (when used as @decorator)

  • input_type – Expected input memory type (default: JAX)

  • output_type – Expected output memory type (default: JAX)

  • oom_recovery – Enable automatic OOM recovery (default: True)

  • contract – Optional validation function for outputs

Returns:

Decorated function with memory type metadata and dtype preservation

Dtype Conversion

class DtypeConversion(value)[source]

Data type conversion modes for all memory type functions.

PRESERVE_INPUT = 'preserve'
NATIVE_OUTPUT = 'native'
UINT8 = 'uint8'
UINT16 = 'uint16'
INT16 = 'int16'
INT32 = 'int32'
FLOAT32 = 'float32'
FLOAT64 = 'float64'
property numpy_dtype

Get the corresponding numpy dtype.

Examples

Basic Usage

from arraybridge import torch
import numpy as np

@torch(input_type='numpy', output_type='torch')
def process_data(data):
    return data * 2

result = process_data(np.array([1, 2, 3]))

With OOM Recovery

from arraybridge import cupy

@cupy(oom_recovery=True)
def gpu_operation(data):
    return data @ data.T

Dtype Preservation

from arraybridge import torch
from arraybridge.decorators import DtypeConversion

@torch()
def process_image(image, dtype_conversion=DtypeConversion.PRESERVE_INPUT):
    # Automatically preserves input dtype
    return image * 2