From 0b775c681d2101bc13fb30135b493b476f635f08 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Tue, 6 Oct 2026 18:00:18 -0400 Subject: [PATCH 1/2] Keep image geometry and allocation on the reference device --- README.md | 14 + src/arraybridge/array_operations.py | 702 +++++++++++++++------------- src/arraybridge/types.py | 45 +- tests/test_native_geometry.py | 65 +++ 4 files changed, 495 insertions(+), 331 deletions(-) create mode 100644 tests/test_native_geometry.py diff --git a/README.md b/README.md index 1db61c4..a0baf12 100644 --- a/README.md +++ b/README.md @@ -71,3 +71,17 @@ pip install "arraybridge[cupy]" ``` Documentation: + +### 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. diff --git a/src/arraybridge/array_operations.py b/src/arraybridge/array_operations.py index 1d359ae..b435926 100644 --- a/src/arraybridge/array_operations.py +++ b/src/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() diff --git a/src/arraybridge/types.py b/src/arraybridge/types.py index aec7895..e7c68c2 100644 --- a/src/arraybridge/types.py +++ b/src/arraybridge/types.py @@ -12,7 +12,7 @@ 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 @@ def astype( """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 @@ def logical_and( """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, ...]: diff --git a/tests/test_native_geometry.py b/tests/test_native_geometry.py new file mode 100644 index 0000000..8bcfc17 --- /dev/null +++ b/tests/test_native_geometry.py @@ -0,0 +1,65 @@ +"""Native geometry must preserve pixels, sharing and the reference device.""" + +import numpy as np +import pytest + +from arraybridge.array_operations import PYCLESPERANTO_OPERATIONS +from arraybridge.types import MemoryType + + +def test_numpy_views_and_independent_reference_allocation(): + source = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] + shaped = MemoryType.NUMPY.reshape(source, (3, 1, 4)) + broadcast = MemoryType.NUMPY.broadcast_to(shaped, (3, 2, 4)) + assert np.shares_memory(shaped, source) + assert np.shares_memory(broadcast, source) + np.testing.assert_array_equal(broadcast, np.broadcast_to(source[:, None], (3, 2, 4))) + ones = MemoryType.NUMPY.ones_like(source, shape=(2, 4), dtype=bool) + assert ones.dtype == np.bool_ + assert not np.shares_memory(ones, source) + np.testing.assert_array_equal(ones, np.ones((2, 4), dtype=bool)) + + +def test_numpy_geometry_rejects_invalid_shapes(): + source = np.arange(6).reshape(2, 3) + with pytest.raises(ValueError): + MemoryType.NUMPY.reshape(source, (7,)) + with pytest.raises(ValueError): + MemoryType.NUMPY.broadcast_to(source, (2, 4)) + + +def test_cupy_geometry_and_allocation_stay_native(): + cp = pytest.importorskip("cupy") + if not MemoryType.CUPY.available_device_ids(cp): + pytest.skip("No CUDA device") + source = cp.asarray(np.arange(12, dtype=np.float32).reshape(3, 4))[:, ::-1] + shaped = MemoryType.CUPY.reshape(source, (3, 1, 4)) + broadcast = MemoryType.CUPY.broadcast_to(shaped, (3, 2, 4)) + assert shaped.data.ptr == source.data.ptr + assert broadcast.data.ptr == source.data.ptr + assert broadcast.device.id == source.device.id + np.testing.assert_array_equal( + broadcast.get(), np.broadcast_to(source.get()[:, None], (3, 2, 4)) + ) + ones = MemoryType.CUPY.ones_like(source, shape=(2, 4), dtype=bool) + assert ones.device.id == source.device.id + assert ones.dtype == cp.bool_ + np.testing.assert_array_equal(ones.get(), np.ones((2, 4), dtype=bool)) + + +def test_pycles_geometry_rejects_host_fallback(): + class DeviceOnlyArray: + shape = (2, 3) + + def __array__(self): + pytest.fail("Native geometry must never download pixels") + + source = DeviceOnlyArray() + assert PYCLESPERANTO_OPERATIONS.reshape(source, (2, 3), None) is source + assert PYCLESPERANTO_OPERATIONS.broadcast_to(source, (2, 3), None) is source + with pytest.raises(NotImplementedError, match="reshape"): + PYCLESPERANTO_OPERATIONS.reshape(source, (1, 2, 3), None) + with pytest.raises(NotImplementedError, match="broadcasting"): + PYCLESPERANTO_OPERATIONS.broadcast_to(source, (4, 2, 3), None) + with pytest.raises(NotImplementedError, match="dtype"): + PYCLESPERANTO_OPERATIONS.ones_like(source, (2, 3), bool, None) From b8a2eabf28697905ff325f9217137fcba14d52b4 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Tue, 6 Oct 2026 18:04:27 -0400 Subject: [PATCH 2/2] Prepare native geometry release 0.3.7 and repair release documentation gates --- docs/source/index.rst | 2 ++ docs/source/release_035_validation.rst | 2 +- docs/source/release_036_validation.rst | 2 +- src/arraybridge/__init__.py | 2 +- tests/test_callable_dtype_default.py | 4 +++- 5 files changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/source/index.rst b/docs/source/index.rst index ab30191..575ebf3 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -20,6 +20,8 @@ applications. examples/index contributing ci-cd + release_035_validation + release_036_validation Boundary rule ------------- diff --git a/docs/source/release_035_validation.rst b/docs/source/release_035_validation.rst index b58bfbc..9909306 100644 --- a/docs/source/release_035_validation.rst +++ b/docs/source/release_035_validation.rst @@ -1,5 +1,5 @@ ArrayBridge 0.3.5 release preparation -=================================== +===================================== Parent integration owner; base ea3f2a4cc91c4810d12343f58f85c1195e1a41e6. The reviewed callable-native dtype default and durable decorator-context diff --git a/docs/source/release_036_validation.rst b/docs/source/release_036_validation.rst index 92de2d3..d6541c4 100644 --- a/docs/source/release_036_validation.rst +++ b/docs/source/release_036_validation.rst @@ -1,5 +1,5 @@ ArrayBridge typed thread-local release correction -================================================ +================================================= Parent integration owner. Actual failed original 0.3.5 publisher: https://github.com/OpenHCSDev/ArrayBridge/actions/runs/36946585391. diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index 209e57d..5985b4c 100644 --- a/src/arraybridge/__init__.py +++ b/src/arraybridge/__init__.py @@ -10,7 +10,7 @@ NumPy/numcodecs import cost. """ -__version__ = "0.3.6" +__version__ = "0.3.7" _LAZY_EXPORTS: dict[str, str] = { "MemoryType": ".types", diff --git a/tests/test_callable_dtype_default.py b/tests/test_callable_dtype_default.py index e3bec06..c2fb14e 100644 --- a/tests/test_callable_dtype_default.py +++ b/tests/test_callable_dtype_default.py @@ -1,7 +1,7 @@ """Callable defaults use the existing typed dtype policy and real conversion.""" -from dataclasses import dataclass import inspect +from dataclasses import dataclass import numpy as np import pytest @@ -10,6 +10,8 @@ DtypeConversion, DtypeConversionConfig, PreserveInputDtypeConfig, +) +from arraybridge.decorators import ( numpy as numpy_func, )