arraybridge 0.3.4__py3-none-any.whl → 0.3.7__py3-none-any.whl

arraybridge/__init__.py

@@ -10,7 +10,7 @@ attribute access instead of at import time. Declaration-only consumers
 NumPy/numcodecs import cost.
 """
 
-__version__ = "0.3.4"
+__version__ = "0.3.7"
 
 _LAZY_EXPORTS: dict[str, str] = {
     "MemoryType": ".types",

arraybridge/array_operations.py

@@ -1,19 +1,11 @@
-"""Typed array-operation leaves carried by ``MemoryType`` declarations."""
+"""Native array operations carried by MemoryType declarations."""
 
-from collections.abc import Callable, Sequence
-from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from typing import Any
 
 import numpy as np
 
-ToNumpy = Callable[[Any, Any], Any]
-FromNumpy = Callable[[Any, Any, int], Any]
-StackArrays = Callable[[Sequence[Any], Any], Any]
-ScaleDtype = Callable[[Any, Any, Any], Any]
-DtypeName = Callable[[Any], str]
-CastArray = Callable[[Any, Any, Any], Any]
-LogicalAnd = Callable[[Any, Any, Any], Any]
-
 _SCALING_RANGES: dict[str, float | tuple[float, float]] = {
     "uint8": 255.0,
     "uint16": 65535.0,
@@ -64,103 +56,6 @@ def _clamp_bounds(target_dtype: Any) -> tuple[float, float] | None:
     return 0, range_info
 
 
-def _identity_to_numpy(data: Any, module: Any) -> Any:
-    del module
-    return data
-
-
-def _identity_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    del module, device_id
-    return data
-
-
-def _numpy_stack(values: Sequence[Any], module: Any) -> Any:
-    return module.stack(values, axis=0)
-
-
-def _numpy_cast(data: Any, dtype: Any, module: Any) -> Any:
-    del module
-    return data.astype(dtype, copy=False)
-
-
-def _module_logical_and(left: Any, right: Any, module: Any) -> Any:
-    return module.logical_and(left, right)
-
-
-def _numpy_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
-        return result
-    if not (
-        module.issubdtype(result.dtype, module.floating)
-        and module.issubdtype(target_dtype, module.integer)
-    ):
-        return result.astype(target_dtype)
-    result_min = result.min()
-    result_max = result.max()
-    if result_max <= result_min:
-        return result.astype(target_dtype)
-    scaled = _scaled_values(result, result_min, result_max, target_dtype)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        scaled = module.clip(scaled, *bounds)
-    return scaled.astype(target_dtype)
-
-
-def _cupy_to_numpy(data: Any, module: Any) -> Any:
-    del module
-    return data.get()
-
-
-def _cupy_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    del device_id
-    return module.array(data)
-
-
-def _cupy_stack(values: Sequence[Any], module: Any) -> Any:
-    return module.stack(values, axis=0)
-
-
-def _cupy_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
-        return result
-    if not (
-        module.issubdtype(result.dtype, module.floating)
-        and not module.issubdtype(target_dtype, module.floating)
-    ):
-        return result.astype(target_dtype)
-    result_min = module.min(result)
-    result_max = module.max(result)
-    if result_max <= result_min:
-        return result.astype(target_dtype)
-    scaled = _scaled_values(result, result_min, result_max, target_dtype)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        scaled = module.clip(scaled, *bounds)
-    return scaled.astype(target_dtype)
-
-
-def _torch_to_numpy(data: Any, module: Any) -> Any:
-    del module
-    return data.cpu().numpy()
-
-
-def _torch_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    host_data = (
-        np.ascontiguousarray(data)
-        if any(stride < 0 for stride in getattr(data, "strides", ()))
-        else data
-    )
-    return module.from_numpy(host_data).to(f"cuda:{device_id}")
-
-
-def _torch_stack(values: Sequence[Any], module: Any) -> Any:
-    return module.stack(tuple(values), dim=0)
-
-
-def _torch_cast(data: Any, dtype: Any, module: Any) -> Any:
-    return data.to(dtype=_mapped_dtype(dtype, module))
-
-
 def _mapped_dtype(target_dtype: Any, module: Any) -> Any:
     try:
         dtype_name = np.dtype(target_dtype).name
@@ -173,224 +68,383 @@ def _mapped_dtype(target_dtype: Any, module: Any) -> Any:
     return mapped
 
 
-def _torch_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
+class ArrayOperations(ABC):
+    """Native array semantics selected by the existing MemoryType declaration."""
+
+    @staticmethod
+    @abstractmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        """Project pixels to the explicit host boundary."""
+
+    @staticmethod
+    @abstractmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        """Admit host pixels on the selected framework device."""
+
+    @staticmethod
+    @abstractmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        """Apply the framework's existing intensity conversion law."""
+
+    @staticmethod
+    def dtype_name(dtype: Any) -> str:
+        return _numpy_dtype_name(dtype)
+
+    @staticmethod
+    def stack(values: Sequence[Any], module: Any) -> Any:
+        return module.stack(tuple(values), axis=0)
+
+    @staticmethod
+    def cast(data: Any, dtype: Any, module: Any) -> Any:
+        return data.astype(dtype, copy=False)
+
+    @staticmethod
+    def logical_and(left: Any, right: Any, module: Any) -> Any:
+        return module.logical_and(left, right)
+
+    @staticmethod
+    def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        return data.reshape(shape)
+
+    @staticmethod
+    def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        return module.broadcast_to(data, shape)
+
+    @staticmethod
+    def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
+        return module.ones(shape, dtype=dtype)
+
+
+class NumpyArrayOperations(ArrayOperations):
+    """Numpy native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        del module
+        return data
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        del module, device_id
+        return data
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        if not (
+            module.issubdtype(result.dtype, module.floating)
+            and module.issubdtype(target_dtype, module.integer)
+        ):
+            return result.astype(target_dtype)
+        result_min = result.min()
+        result_max = result.max()
+        if result_max <= result_min:
+            return result.astype(target_dtype)
+        scaled = _scaled_values(result, result_min, result_max, target_dtype)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            scaled = module.clip(scaled, *bounds)
+        return scaled.astype(target_dtype)
+
+
+class CupyArrayOperations(ArrayOperations):
+    """Cupy native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        del module
+        return data.get()
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        del device_id
+        return module.array(data)
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        if not (
+            module.issubdtype(result.dtype, module.floating)
+            and not module.issubdtype(target_dtype, module.floating)
+        ):
+            return result.astype(target_dtype)
+        result_min = module.min(result)
+        result_max = module.max(result)
+        if result_max <= result_min:
+            return result.astype(target_dtype)
+        scaled = _scaled_values(result, result_min, result_max, target_dtype)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            scaled = module.clip(scaled, *bounds)
+        return scaled.astype(target_dtype)
+
+
+class TorchArrayOperations(ArrayOperations):
+    """Torch native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        del module
+        return data.cpu().numpy()
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        host_data = (
+            np.ascontiguousarray(data)
+            if any(stride < 0 for stride in getattr(data, "strides", ()))
+            else data
+        )
+        return module.from_numpy(host_data).to(f"cuda:{device_id}")
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        mapped = _mapped_dtype(target_dtype, module)
+        floats = (module.float16, module.float32, module.float64)
+        if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
+            return result.to(mapped)
+        result_min = result.min()
+        result_max = result.max()
+        if result_max <= result_min:
+            return result.to(mapped)
+        scaled = _scaled_values(result, result_min, result_max, target_dtype)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            scaled = module.clamp(scaled, min=bounds[0], max=bounds[1])
+        return scaled.to(mapped)
+
+    @staticmethod
+    def stack(values: Sequence[Any], module: Any) -> Any:
+        return module.stack(tuple(values), dim=0)
+
+    @staticmethod
+    def cast(data: Any, dtype: Any, module: Any) -> Any:
+        return data.to(dtype=_mapped_dtype(dtype, module))
+
+    @staticmethod
+    def dtype_name(dtype: Any) -> str:
+        return str(dtype).rsplit(".", maxsplit=1)[-1]
+
+    @staticmethod
+    def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        return data.expand(shape)
+
+    @staticmethod
+    def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
+        return module.ones(shape, dtype=_mapped_dtype(dtype, module), device=reference.device)
+
+
+class TensorflowArrayOperations(ArrayOperations):
+    """Tensorflow native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        del module
+        return data.numpy()
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        del device_id
+        return module.convert_to_tensor(data)
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        mapped = _mapped_dtype(target_dtype, module)
+        floats = (module.float16, module.float32, module.float64)
+        if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
+            return module.cast(result, mapped)
+        result_min = module.reduce_min(result)
+        result_max = module.reduce_max(result)
+        if result_max <= result_min:
+            return module.cast(result, mapped)
+        scaled = _scaled_values(result, result_min, result_max, target_dtype)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            scaled = module.clip_by_value(scaled, *bounds)
+        return module.cast(scaled, mapped)
+
+    @staticmethod
+    def cast(data: Any, dtype: Any, module: Any) -> Any:
+        return module.cast(data, _mapped_dtype(dtype, module))
+
+    @staticmethod
+    def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        return module.reshape(data, shape)
+
+    @staticmethod
+    def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
+        with module.device(reference.device):
+            return module.ones(shape, dtype=_mapped_dtype(dtype, module))
+
+
+class JaxArrayOperations(ArrayOperations):
+    """Jax native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        del module
+        return np.asarray(data)
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        devices = tuple(device for device in module.devices() if device.platform == "gpu")
+        return module.device_put(data, devices[device_id])
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        if np.dtype(target_dtype) == np.dtype(np.float64):
+            x64_enabled = getattr(module.config, "x64_enabled", None)
+            if x64_enabled is None:
+                x64_enabled = module.config.read("jax_enable_x64")
+            if not x64_enabled:
+                raise ValueError(
+                    "JAX float64 output requires x64 mode; set JAX_ENABLE_X64=true before import"
+                )
+        jnp = module.numpy
+        mapped = _mapped_dtype(target_dtype, jnp)
+        floats = (jnp.float16, jnp.float32, jnp.float64)
+        if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
+            return result.astype(mapped)
+        result_min = jnp.min(result)
+        result_max = jnp.max(result)
+        if result_max <= result_min:
+            return result.astype(mapped)
+        scaled = _scaled_values(result, result_min, result_max, target_dtype)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            scaled = jnp.clip(scaled, *bounds)
+        return scaled.astype(mapped)
+
+    @staticmethod
+    def stack(values: Sequence[Any], module: Any) -> Any:
+        return module.numpy.stack(tuple(values), axis=0)
+
+    @staticmethod
+    def cast(data: Any, dtype: Any, module: Any) -> Any:
+        return data.astype(_mapped_dtype(dtype, module.numpy))
+
+    @staticmethod
+    def logical_and(left: Any, right: Any, module: Any) -> Any:
+        return module.numpy.logical_and(left, right)
+
+    @staticmethod
+    def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        return module.numpy.broadcast_to(data, shape)
+
+    @staticmethod
+    def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
+        return module.numpy.ones(shape, dtype=dtype, device=reference.device)
+
+
+class PyclesperantoArrayOperations(ArrayOperations):
+    """Pyclesperanto native operation leaves."""
+
+    @staticmethod
+    def to_numpy(data: Any, module: Any) -> Any:
+        return module.pull(data)
+
+    @staticmethod
+    def from_numpy(data: Any, module: Any, device_id: int) -> Any:
+        del device_id
+        return module.push(data)
+
+    @staticmethod
+    def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
+        if not hasattr(result, "dtype"):
+            return result
+        target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer)
+        if not (np.issubdtype(result.dtype, np.floating) and target_is_int):
+            return module.push(module.pull(result).astype(target_dtype))
+        result_min = float(module.minimum_of_all_pixels(result))
+        result_max = float(module.maximum_of_all_pixels(result))
+        if result_max <= result_min:
+            return module.push(module.pull(result).astype(target_dtype))
+        normalized = module.subtract_image_from_scalar(result, scalar=result_min)
+        normalized = module.multiply_image_and_scalar(
+            normalized,
+            scalar=1.0 / (result_max - result_min),
+        )
+        range_info = _SCALING_RANGES.get(_dtype_name(target_dtype))
+        if isinstance(range_info, tuple):
+            scale, offset = range_info
+            scaled = module.multiply_image_and_scalar(normalized, scalar=scale)
+            scaled = module.subtract_image_from_scalar(scaled, scalar=offset)
+        elif range_info is not None:
+            scaled = module.multiply_image_and_scalar(normalized, scalar=range_info)
+        else:
+            scaled = normalized
+        host_values = module.pull(scaled)
+        bounds = _clamp_bounds(target_dtype)
+        if bounds is not None:
+            host_values = np.clip(host_values, *bounds)
+        return module.push(host_values.astype(target_dtype))
+
+    @staticmethod
+    def stack(values: Sequence[Any], module: Any) -> Any:
+        if not values:
+            raise ValueError("Cannot stack an empty pyclesperanto sequence")
+        if len(values) == 1:
+            source = values[0]
+            result = module.create((1, *source.shape), dtype=source.dtype)
+            return module.copy_slice(source, result, 0)
+        result = values[0]
+        for value in values[1:]:
+            result = module.concatenate_along_z(result, value)
         return result
-    mapped = _mapped_dtype(target_dtype, module)
-    floats = (module.float16, module.float32, module.float64)
-    if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
-        return result.to(mapped)
-    result_min = result.min()
-    result_max = result.max()
-    if result_max <= result_min:
-        return result.to(mapped)
-    scaled = _scaled_values(result, result_min, result_max, target_dtype)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        scaled = module.clamp(scaled, min=bounds[0], max=bounds[1])
-    return scaled.to(mapped)
-
-
-def _tensorflow_to_numpy(data: Any, module: Any) -> Any:
-    del module
-    return data.numpy()
-
-
-def _tensorflow_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    del device_id
-    return module.convert_to_tensor(data)
 
-
-def _tensorflow_stack(values: Sequence[Any], module: Any) -> Any:
-    return module.stack(tuple(values), axis=0)
-
-
-def _tensorflow_cast(data: Any, dtype: Any, module: Any) -> Any:
-    return module.cast(data, _mapped_dtype(dtype, module))
-
-
-def _tensorflow_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
-        return result
-    mapped = _mapped_dtype(target_dtype, module)
-    floats = (module.float16, module.float32, module.float64)
-    if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
-        return module.cast(result, mapped)
-    result_min = module.reduce_min(result)
-    result_max = module.reduce_max(result)
-    if result_max <= result_min:
-        return module.cast(result, mapped)
-    scaled = _scaled_values(result, result_min, result_max, target_dtype)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        scaled = module.clip_by_value(scaled, *bounds)
-    return module.cast(scaled, mapped)
-
-
-def _jax_to_numpy(data: Any, module: Any) -> Any:
-    del module
-    return np.asarray(data)
-
-
-def _jax_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    devices = tuple(device for device in module.devices() if device.platform == "gpu")
-    return module.device_put(data, devices[device_id])
-
-
-def _jax_stack(values: Sequence[Any], module: Any) -> Any:
-    return module.numpy.stack(tuple(values), axis=0)
-
-
-def _jax_cast(data: Any, dtype: Any, module: Any) -> Any:
-    return data.astype(_mapped_dtype(dtype, module.numpy))
-
-
-def _jax_logical_and(left: Any, right: Any, module: Any) -> Any:
-    return module.numpy.logical_and(left, right)
-
-
-def _jax_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
-        return result
-    if np.dtype(target_dtype) == np.dtype(np.float64):
-        x64_enabled = getattr(module.config, "x64_enabled", None)
-        if x64_enabled is None:
-            x64_enabled = module.config.read("jax_enable_x64")
-        if not x64_enabled:
-            raise ValueError(
-                "JAX float64 output requires x64 mode; set JAX_ENABLE_X64=true before import"
+    @staticmethod
+    def cast(data: Any, dtype: Any, module: Any) -> Any:
+        return module.push(module.pull(data).astype(dtype, copy=False))
+
+    @staticmethod
+    def logical_and(left: Any, right: Any, module: Any) -> Any:
+        return module.push(np.logical_and(module.pull(left), module.pull(right)))
+
+    @staticmethod
+    def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        if tuple(data.shape) == shape:
+            return data
+        raise NotImplementedError(
+            "pyclesperanto does not provide native array reshape; its reshape downloads pixels"
+        )
+
+    @staticmethod
+    def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any:
+        if tuple(data.shape) == shape:
+            return data
+        raise NotImplementedError("pyclesperanto does not provide native array broadcasting")
+
+    @staticmethod
+    def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
+        if np.dtype(dtype).name not in {
+            "float32",
+            "int8",
+            "int16",
+            "int32",
+            "uint8",
+            "uint16",
+            "uint32",
+        }:
+            raise NotImplementedError(f"pyclesperanto cannot allocate native dtype {dtype!r}")
+        if not 1 <= len(shape) <= 3:
+            raise NotImplementedError(
+                "pyclesperanto native allocation requires one to three dimensions"
             )
-    jnp = module.numpy
-    mapped = _mapped_dtype(target_dtype, jnp)
-    floats = (jnp.float16, jnp.float32, jnp.float64)
-    if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
-        return result.astype(mapped)
-    result_min = jnp.min(result)
-    result_max = jnp.max(result)
-    if result_max <= result_min:
-        return result.astype(mapped)
-    scaled = _scaled_values(result, result_min, result_max, target_dtype)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        scaled = jnp.clip(scaled, *bounds)
-    return scaled.astype(mapped)
-
-
-def _pyclesperanto_to_numpy(data: Any, module: Any) -> Any:
-    return module.pull(data)
-
-
-def _pyclesperanto_from_numpy(data: Any, module: Any, device_id: int) -> Any:
-    del device_id
-    return module.push(data)
-
-
-def _pyclesperanto_stack(values: Sequence[Any], module: Any) -> Any:
-    if not values:
-        raise ValueError("Cannot stack an empty pyclesperanto sequence")
-    if len(values) == 1:
-        source = values[0]
-        result = module.create((1, *source.shape), dtype=source.dtype)
-        return module.copy_slice(source, result, 0)
-    result = values[0]
-    for value in values[1:]:
-        result = module.concatenate_along_z(result, value)
-    return result
-
-
-def _pyclesperanto_cast(data: Any, dtype: Any, module: Any) -> Any:
-    return module.push(module.pull(data).astype(dtype, copy=False))
-
-
-def _pyclesperanto_logical_and(left: Any, right: Any, module: Any) -> Any:
-    return module.push(np.logical_and(module.pull(left), module.pull(right)))
+        result = module.create(shape, dtype=dtype, device=reference.device)
+        module.set(result, scalar=1, device=reference.device)
+        return result
 
 
-def _pyclesperanto_scale(result: Any, target_dtype: Any, module: Any) -> Any:
-    if not hasattr(result, "dtype"):
-        return result
-    target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer)
-    if not (np.issubdtype(result.dtype, np.floating) and target_is_int):
-        return module.push(module.pull(result).astype(target_dtype))
-    result_min = float(module.minimum_of_all_pixels(result))
-    result_max = float(module.maximum_of_all_pixels(result))
-    if result_max <= result_min:
-        return module.push(module.pull(result).astype(target_dtype))
-    normalized = module.subtract_image_from_scalar(result, scalar=result_min)
-    normalized = module.multiply_image_and_scalar(
-        normalized,
-        scalar=1.0 / (result_max - result_min),
-    )
-    range_info = _SCALING_RANGES.get(_dtype_name(target_dtype))
-    if isinstance(range_info, tuple):
-        scale, offset = range_info
-        scaled = module.multiply_image_and_scalar(normalized, scalar=scale)
-        scaled = module.subtract_image_from_scalar(scaled, scalar=offset)
-    elif range_info is not None:
-        scaled = module.multiply_image_and_scalar(normalized, scalar=range_info)
-    else:
-        scaled = normalized
-    host_values = module.pull(scaled)
-    bounds = _clamp_bounds(target_dtype)
-    if bounds is not None:
-        host_values = np.clip(host_values, *bounds)
-    return module.push(host_values.astype(target_dtype))
-
-
-@dataclass(frozen=True, slots=True)
-class ArrayOperations:
-    """Framework-specific array leaves referenced by one declaration."""
-
-    to_numpy: ToNumpy
-    from_numpy: FromNumpy
-    stack: StackArrays
-    scale_dtype: ScaleDtype
-    dtype_name: DtypeName = _numpy_dtype_name
-    cast: CastArray = _numpy_cast
-    logical_and: LogicalAnd = _module_logical_and
-
-
-NUMPY_OPERATIONS = ArrayOperations(
-    to_numpy=_identity_to_numpy,
-    from_numpy=_identity_from_numpy,
-    stack=_numpy_stack,
-    scale_dtype=_numpy_scale,
-)
-CUPY_OPERATIONS = ArrayOperations(
-    to_numpy=_cupy_to_numpy,
-    from_numpy=_cupy_from_numpy,
-    stack=_cupy_stack,
-    scale_dtype=_cupy_scale,
-)
-TORCH_OPERATIONS = ArrayOperations(
-    to_numpy=_torch_to_numpy,
-    from_numpy=_torch_from_numpy,
-    stack=_torch_stack,
-    scale_dtype=_torch_scale,
-    dtype_name=_torch_dtype_name,
-    cast=_torch_cast,
-)
-TENSORFLOW_OPERATIONS = ArrayOperations(
-    to_numpy=_tensorflow_to_numpy,
-    from_numpy=_tensorflow_from_numpy,
-    stack=_tensorflow_stack,
-    scale_dtype=_tensorflow_scale,
-    dtype_name=_tensorflow_dtype_name,
-    cast=_tensorflow_cast,
-)
-JAX_OPERATIONS = ArrayOperations(
-    to_numpy=_jax_to_numpy,
-    from_numpy=_jax_from_numpy,
-    stack=_jax_stack,
-    scale_dtype=_jax_scale,
-    cast=_jax_cast,
-    logical_and=_jax_logical_and,
-)
-PYCLESPERANTO_OPERATIONS = ArrayOperations(
-    to_numpy=_pyclesperanto_to_numpy,
-    from_numpy=_pyclesperanto_from_numpy,
-    stack=_pyclesperanto_stack,
-    scale_dtype=_pyclesperanto_scale,
-    cast=_pyclesperanto_cast,
-    logical_and=_pyclesperanto_logical_and,
-)
+NUMPY_OPERATIONS = NumpyArrayOperations()
+CUPY_OPERATIONS = CupyArrayOperations()
+TORCH_OPERATIONS = TorchArrayOperations()
+TENSORFLOW_OPERATIONS = TensorflowArrayOperations()
+JAX_OPERATIONS = JaxArrayOperations()
+PYCLESPERANTO_OPERATIONS = PyclesperantoArrayOperations()

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

@@ -12,7 +12,7 @@ import importlib.util
 import logging
 import os
 import sys
-from collections.abc import Callable, Iterator, Mapping, MutableMapping
+from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from enum import Enum
@@ -834,11 +834,44 @@ class MemoryType(_MemoryTypeFields, Enum):
         """Cast one array through this framework member's operation leaf."""
 
         framework = module if module is not None else self.import_module()
-        device_id = self.device_id_of(data, framework)
-        scope = nullcontext() if device_id is None else self.device_scope(device_id, framework)
-        with scope:
+        with self._array_device_scope(data, framework):
             return self._operations.cast(data, dtype, framework)
 
+    def _array_device_scope(self, reference: Any, module: Any) -> AbstractContextManager[None]:
+        """Derive every native operation scope from its actual input array."""
+        device_id = self.device_id_of(reference, module)
+        return nullcontext() if device_id is None else self.device_scope(device_id, module)
+
+    def reshape(self, data: Any, shape: Sequence[int], module: Any | None = None) -> Any:
+        """Reshape on the input device without an implicit host projection."""
+        framework = module if module is not None else self.import_module()
+        with self._array_device_scope(data, framework):
+            return self._operations.reshape(data, tuple(shape), framework)
+
+    def broadcast_to(self, data: Any, shape: Sequence[int], module: Any | None = None) -> Any:
+        """Broadcast on the input device without an implicit host projection."""
+        framework = module if module is not None else self.import_module()
+        with self._array_device_scope(data, framework):
+            return self._operations.broadcast_to(data, tuple(shape), framework)
+
+    def ones_like(
+        self,
+        reference: Any,
+        *,
+        shape: Sequence[int] | None = None,
+        dtype: Any = bool,
+        module: Any | None = None,
+    ) -> Any:
+        """Allocate ones on the supplied reference's framework-local device."""
+        framework = module if module is not None else self.import_module()
+        with self._array_device_scope(reference, framework):
+            return self._operations.ones_like(
+                reference,
+                tuple(reference.shape if shape is None else shape),
+                dtype,
+                framework,
+            )
+
     def logical_and(
         self,
         left: Any,
@@ -848,9 +881,7 @@ class MemoryType(_MemoryTypeFields, Enum):
         """Intersect two arrays through this framework member's operation leaf."""
 
         framework = module if module is not None else self.import_module()
-        device_id = self.device_id_of(left, framework)
-        scope = nullcontext() if device_id is None else self.device_scope(device_id, framework)
-        with scope:
+        with self._array_device_scope(left, framework):
             return self._operations.logical_and(left, right, framework)
 
     def available_device_ids(self, module: Any | None = None) -> tuple[int, ...]:

{arraybridge-0.3.4.dist-info → arraybridge-0.3.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arraybridge
-Version: 0.3.4
+Version: 0.3.7
 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'
@@ -142,3 +143,17 @@ pip install "arraybridge[cupy]"
 ```
 
 Documentation: <https://arraybridge.readthedocs.io>
+
+### Native array geometry
+
+`MemoryType.reshape(array, shape)` and `MemoryType.broadcast_to(array, shape)`
+operate in the array's framework and device. `MemoryType.ones_like(reference,
+shape=..., dtype=...)` allocates on the reference's device; omitting `shape`
+retains its shape. Reshape and broadcast preserve views where the backend can
+represent them. Broadcasting is not a promise of a writable, independent buffer.
+
+These operations never silently project pixels to NumPy. Pyclesperanto cannot
+provide native changed-shape reshape or broadcasting, so those requests raise
+`NotImplementedError`. Its native ones allocation supports one to three
+dimensions and its supported integer/float32 dtypes; boolean allocation is
+explicitly unsupported. Existing conversion and disk boundaries remain separate.

{arraybridge-0.3.4.dist-info → arraybridge-0.3.7.dist-info}/RECORD

@@ -1,10 +1,10 @@
-arraybridge/__init__.py,sha256=a27ES1ItSBsT77vP3f12ILNu0MzrvQfcpCBVDXwi-Sg,3184
+arraybridge/__init__.py,sha256=lmu9luoyMe9tpvuDlxSvwsYhRcFKkLKP7Smm93RndS8,3184
 arraybridge/array_geometry.py,sha256=CopfGmFVKwV5l-lOTImhm1YVh18weUng0ReEF58qlPI,1689
-arraybridge/array_operations.py,sha256=fp69o9BYcCzfhmLcWrH68SVGSfPL3nb6Cty-Amd_oQM,13082
+arraybridge/array_operations.py,sha256=3cT17uSpq9tQPj_HkxHfR8_q6C8tzygI2HiM0yuAaAk,16381
 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=e5R2DVflQAjXrfTk033KRdiiwoip3IQAvgDCYUjPdiQ,35807
+arraybridge/types.py,sha256=_RrqtA2q4ocy1huxLIDiDhoZySCQpNhorV5_7HY0S-Y,37287
 arraybridge/utils.py,sha256=RqThtYScPEOEQG491dA8b_VvUTjgcFhwLbh0erQu5Ns,4833
-arraybridge-0.3.4.dist-info/METADATA,sha256=jq9-ULcNwbZo7cMtZlhhqj4F2QlV-8PrpdM1Tv0tLCE,5275
-arraybridge-0.3.4.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
-arraybridge-0.3.4.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
-arraybridge-0.3.4.dist-info/RECORD,,
+arraybridge-0.3.7.dist-info/METADATA,sha256=Rpr8OMPyw7t9BXgM3EWf2jAm03vxUdVIPeGN3L29FmU,6124
+arraybridge-0.3.7.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+arraybridge-0.3.7.dist-info/licenses/LICENSE,sha256=xagEoeTAj1WT64RmyR3E6HH-eTGdgXN6gqPMUUt7L_Y,1070
+arraybridge-0.3.7.dist-info/RECORD,,