From fdcd2c64d60fd6b421cd6865ab57bf33d7b5c199 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Sat, 10 Oct 2026 21:32:09 -0400 Subject: [PATCH 1/2] ArrayOperations is a registered family with for_memory(); leaves hold the primitives that differ ArrayOperations leaves register by memory_type, and ArrayOperations.for_memory(MemoryType.X) returns the leaf. MemoryType finds its leaf through the family; the operations-name field and the six module-level *_OPERATIONS singletons are deleted. Leaves gain an instance namespace for array-library-neutral code: module (imported on first use), accepts, astype, max, min, mean, sum, prepend_axis, linspace, ones, assign, outer and floor. Each leaf overrides only what its library does differently (torch/JAX/TensorFlow float32 means, torch uint16 reductions, functional updates for JAX and TensorFlow, pyclesperanto z-projections). torch and TensorFlow cast also accepts native dtypes. Version 0.4.0. Co-Authored-By: Claude Opus 5.5 --- ADDING_NEW_FRAMEWORKS.md | 13 +- CONTRIBUTING.md | 4 +- README.md | 24 +++ src/arraybridge/__init__.py | 2 +- src/arraybridge/array_operations.py | 227 ++++++++++++++++++++++++-- src/arraybridge/types.py | 16 +- tests/test_array_operations_family.py | 65 ++++++++ tests/test_native_geometry.py | 4 +- tests/test_stack_utils.py | 9 +- 9 files changed, 325 insertions(+), 39 deletions(-) create mode 100644 tests/test_array_operations_family.py diff --git a/ADDING_NEW_FRAMEWORKS.md b/ADDING_NEW_FRAMEWORKS.md index db3daa3..82a03c7 100644 --- a/ADDING_NEW_FRAMEWORKS.md +++ b/ADDING_NEW_FRAMEWORKS.md @@ -6,13 +6,17 @@ own framework registry. ## 1. Define typed leaves -Add an `ArrayOperations` bundle in -`src/arraybridge/array_operations.py`. It must provide four ordinary callables: +Add an `ArrayOperations` leaf in `src/arraybridge/array_operations.py`: a +subclass declaring `memory_type = ""`. It registers itself, and +`ArrayOperations.for_memory(MemoryType.)` returns it. It must provide: - project an array to NumPy; - create an array from NumPy on the requested framework-local device; - stack prepared 2D arrays into one 3D array; -- apply the framework's dtype-scaling semantics. +- apply the framework's dtype-scaling semantics; +- an override of each namespace primitive (`max`, `mean`, `assign`, + `linspace`, ...) whose library call or numeric rule differs from the + NumPy-compatible default. Define device, stream, DLPack, cleanup, and OOM leaves in `src/arraybridge/types.py`. Use the no-op defaults in `FrameworkRuntime` only @@ -27,8 +31,7 @@ Add one `MemoryType` member carrying: 2. its display name and module aliases; 3. whether it is GPU-backed; 4. any environment defaults that must be set before first import; -5. its `FrameworkRuntime` leaves; -6. its `ArrayOperations` bundle. +5. its `FrameworkRuntime` leaves. The converter classes, target conversion methods, compatibility operation view, dtype-scaling view, cleanup adapters, and framework decorator are generated diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5bd8b90..cd3a403 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -404,8 +404,8 @@ GitHub Actions will automatically: To add support for a new framework: -1. Add typed array-operation leaves in `array_operations.py`. -2. Add one `MemoryType` declaration carrying those operations and its runtime leaves. +1. Add one `ArrayOperations` leaf in `array_operations.py` declaring its `memory_type`. +2. Add one `MemoryType` declaration carrying its runtime leaves. 3. Add detection, conversion-pair, stack, cross-device, and optional-dependency tests. 4. Update the extension and reference documentation. diff --git a/README.md b/README.md index a0baf12..eeb298c 100644 --- a/README.md +++ b/README.md @@ -85,3 +85,27 @@ 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. + +### Writing an operation once for every framework + +`ArrayOperations.for_memory(memory_type)` returns that framework's operations +leaf: one registered subclass per `MemoryType`, holding only the primitives its +library does differently (`max`, `min`, `mean`, `sum`, `astype`, +`prepend_axis`, `linspace`, `ones`, `assign`, `outer`, `floor`, `accepts`). +Code written against the leaf runs natively in each framework: + +```python +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations + + +def max_projection(ops, stack): + return ops.prepend_axis(ops.max(stack, 0)) + + +ops = ArrayOperations.for_memory(MemoryType.TORCH) +``` + +`ops.module` is the framework, imported on first use. Each leaf keeps its +library's numeric rules: torch, JAX and TensorFlow take means in float32, and +pyclesperanto reduces only along the leading axis. diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index eb83370..28b3802 100644 --- a/src/arraybridge/__init__.py +++ b/src/arraybridge/__init__.py @@ -10,7 +10,7 @@ NumPy/numcodecs import cost. """ -__version__ = "0.3.8" +__version__ = "0.4.0" _LAZY_EXPORTS: dict[str, str] = { "MemoryType": ".types", diff --git a/src/arraybridge/array_operations.py b/src/arraybridge/array_operations.py index 90f370c..dc99c91 100644 --- a/src/arraybridge/array_operations.py +++ b/src/arraybridge/array_operations.py @@ -1,10 +1,19 @@ -"""Native array operations carried by MemoryType declarations.""" +"""Native array operations, one family leaf per MemoryType declaration.""" -from abc import ABC, abstractmethod -from collections.abc import Sequence -from typing import Any +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Mapping, Sequence +from functools import cache, cached_property +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np +from metaclass_registry import AutoRegisterMeta + +from arraybridge.array_payload import ArrayPayload + +if TYPE_CHECKING: + from arraybridge.types import MemoryType _SCALING_RANGES: dict[str, float | tuple[float, float]] = { "uint8": 255.0, @@ -68,8 +77,84 @@ def _mapped_dtype(target_dtype: Any, module: Any) -> Any: return mapped -class ArrayOperations(ABC): - """Native array semantics selected by the existing MemoryType declaration.""" +@cache +def _leaf_instance(leaf: type[ArrayOperations]) -> ArrayOperations: + return leaf() + + +class ArrayOperations(metaclass=AutoRegisterMeta): + """Native array semantics, one registered leaf per ``MemoryType`` value. + + Static leaves take the framework module from their caller (``MemoryType``). + Instance primitives are the namespace array-library-neutral code is written + against: each leaf supplies only what its library genuinely does differently. + """ + + __registry_key__ = "memory_type" + # Leaves are declared in this module, without lazy discovery. + __registry__: ClassVar[Mapping[str, type[ArrayOperations]]] = {} + memory_type: ClassVar[str | None] = None + + @classmethod + def for_memory(cls, memory_type: MemoryType) -> ArrayOperations: + """Return the operations leaf declared for ``memory_type``.""" + return _leaf_instance(cls.__registry__[memory_type.value]) + + @cached_property + def module(self) -> Any: + """The framework module, imported on first use.""" + from arraybridge.types import MemoryType + + return MemoryType(self.memory_type).import_module() + + @property + def array_module(self) -> Any: + """The module holding the framework's array functions.""" + return self.module + + # ----- primitives for array-library-neutral operations ----- + + def accepts(self, value: Any) -> bool: + """Whether ``value`` is a native array of this framework.""" + return isinstance(value, self.module.ndarray) + + def astype(self, data: Any, dtype: Any) -> Any: + return self.cast(data, dtype, self.module) + + def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.array_module.max(data, axis=axis) + + def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.array_module.min(data, axis=axis) + + def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.array_module.mean(data, axis=axis) + + def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.array_module.sum(data, axis=axis) + + def prepend_axis(self, data: Any) -> Any: + """Add a leading axis of length one.""" + return self.reshape(data, (1, *data.shape), self.module) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: + return self.array_module.linspace(start, stop, num, endpoint=endpoint) + + def ones(self, shape: int | tuple[int, ...], dtype: Any) -> Any: + """Allocate ones on the framework's default device.""" + return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) + + def assign(self, data: Any, index: slice, values: Any) -> Any: + """Return ``data`` with ``data[index]`` replaced by ``values``.""" + data[index] = values + return data + + def outer(self, left: Any, right: Any) -> Any: + return self.array_module.outer(left, right) + + def floor(self, value: float) -> Any: + """Floor a host scalar in the framework's default precision.""" + return self.array_module.floor(value) @staticmethod @abstractmethod @@ -160,6 +245,11 @@ def normalize_planes( class NumpyArrayOperations(MutableArrayOperations): """Numpy native operation leaves.""" + memory_type = "numpy" + + def accepts(self, value: Any) -> bool: + return isinstance(value, (np.ndarray, ArrayPayload)) + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module @@ -193,6 +283,8 @@ def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: class CupyArrayOperations(MutableArrayOperations): """Cupy native operation leaves.""" + memory_type = "cupy" + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module @@ -226,6 +318,35 @@ def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: class TorchArrayOperations(ArrayOperations): """Torch native operation leaves.""" + memory_type = "torch" + + def accepts(self, value: Any) -> bool: + return isinstance(value, self.module.Tensor) + + def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: + # CUDA has no uint16 reduction; float32 holds every uint16 exactly. + values = data.float() if data.dtype == self.module.uint16 else data + return self.module.amax(values, dim=axis).to(data.dtype) + + def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: + values = data.float() if data.dtype == self.module.uint16 else data + return self.module.amin(values, dim=axis).to(data.dtype) + + def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.mean(data.float(), dim=axis) + + def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.sum(data, dim=axis) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: + float32 = self.module.float32 + if endpoint: + return self.module.linspace(start, stop, num, dtype=float32) + return self.module.linspace(start, stop, num + 1, dtype=float32)[:-1] + + def floor(self, value: float) -> Any: + return self.module.floor(self.module.tensor(value)) + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module @@ -264,7 +385,8 @@ def stack(values: Sequence[Any], module: Any) -> Any: @staticmethod def cast(data: Any, dtype: Any, module: Any) -> Any: - return data.to(dtype=_mapped_dtype(dtype, module)) + native = dtype if isinstance(dtype, module.dtype) else _mapped_dtype(dtype, module) + return data.to(dtype=native) @staticmethod def dtype_name(dtype: Any) -> str: @@ -282,6 +404,40 @@ def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) - class TensorflowArrayOperations(ArrayOperations): """Tensorflow native operation leaves.""" + memory_type = "tensorflow" + + def accepts(self, value: Any) -> bool: + return isinstance(value, self.module.Tensor) + + def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.reduce_max(data, axis=axis) + + def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.reduce_min(data, axis=axis) + + def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.reduce_mean(self.module.cast(data, self.module.float32), axis=axis) + + def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: + return self.module.reduce_sum(data, axis=axis) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> Any: + if endpoint: + return self.module.linspace(float(start), float(stop), num) + return self.module.linspace(float(start), float(stop), num + 1)[:-1] + + def assign(self, data: Any, index: slice, values: Any) -> Any: + rows = self.module.range(*index.indices(int(data.shape[0]))) + return self.module.tensor_scatter_nd_update( + data, self.module.reshape(rows, [-1, 1]), values + ) + + def outer(self, left: Any, right: Any) -> Any: + return self.module.tensordot(left, right, axes=0) + + def floor(self, value: float) -> Any: + return self.module.math.floor(value) + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module @@ -312,7 +468,8 @@ def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any: @staticmethod def cast(data: Any, dtype: Any, module: Any) -> Any: - return module.cast(data, _mapped_dtype(dtype, module)) + native = dtype if isinstance(dtype, module.dtypes.DType) else _mapped_dtype(dtype, module) + return module.cast(data, native) @staticmethod def reshape(data: Any, shape: tuple[int, ...], module: Any) -> Any: @@ -327,6 +484,22 @@ def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) - class JaxArrayOperations(ArrayOperations): """Jax native operation leaves.""" + memory_type = "jax" + + @property + def array_module(self) -> Any: + return self.module.numpy + + def accepts(self, value: Any) -> bool: + return isinstance(value, self.module.numpy.ndarray) + + def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: + jnp = self.module.numpy + return jnp.mean(data.astype(jnp.float32), axis=axis) + + def assign(self, data: Any, index: slice, values: Any) -> Any: + return data.at[index].set(values) + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module @@ -388,6 +561,33 @@ def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) - class PyclesperantoArrayOperations(ArrayOperations): """Pyclesperanto native operation leaves.""" + memory_type = "pyclesperanto" + + def accepts(self, value: Any) -> bool: + # pyclesperanto kernels push any host array they receive. + return hasattr(value, "shape") + + def max(self, data: Any, axis: int | tuple[int, ...]) -> Any: + _require_z_axis(axis) + return self.module.maximum_z_projection(data) + + def min(self, data: Any, axis: int | tuple[int, ...]) -> Any: + _require_z_axis(axis) + return self.module.minimum_z_projection(data) + + def mean(self, data: Any, axis: int | tuple[int, ...]) -> Any: + _require_z_axis(axis) + return self.module.mean_z_projection(data) + + def sum(self, data: Any, axis: int | tuple[int, ...]) -> Any: + _require_z_axis(axis) + return self.module.sum_z_projection(data) + + def prepend_axis(self, data: Any) -> Any: + result = self.module.create((1, *data.shape), dtype=data.dtype) + result[0] = data + return result + @staticmethod def to_numpy(data: Any, module: Any) -> Any: return module.pull(data) @@ -484,9 +684,8 @@ def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) - return result -NUMPY_OPERATIONS = NumpyArrayOperations() -CUPY_OPERATIONS = CupyArrayOperations() -TORCH_OPERATIONS = TorchArrayOperations() -TENSORFLOW_OPERATIONS = TensorflowArrayOperations() -JAX_OPERATIONS = JaxArrayOperations() -PYCLESPERANTO_OPERATIONS = PyclesperantoArrayOperations() +def _require_z_axis(axis: int | tuple[int, ...]) -> None: + if axis != 0: + raise NotImplementedError( + f"pyclesperanto reduces only along the leading (z) axis, not axis {axis!r}" + ) diff --git a/src/arraybridge/types.py b/src/arraybridge/types.py index 7ec4d64..36268fc 100644 --- a/src/arraybridge/types.py +++ b/src/arraybridge/types.py @@ -556,7 +556,6 @@ class _MemoryTypeFields: module_aliases: tuple[str, ...] import_environment: tuple[tuple[str, str], ...] _runtime: FrameworkRuntime - _operations_name: str class MemoryType(_MemoryTypeFields, Enum): @@ -574,7 +573,6 @@ def __new__( module_aliases, import_environment, runtime, - operations, ) = declaration member = object.__new__(cls) member._value_ = value @@ -584,17 +582,15 @@ def __new__( 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_name = cast(str, operations) return member @property def _operations(self) -> ArrayOperations: - """Resolve this declaration's operations registry on first access.""" + """Resolve this declaration's operations leaf on first access.""" - from arraybridge import array_operations + from arraybridge.array_operations import ArrayOperations - operations: ArrayOperations = getattr(array_operations, self._operations_name) - return operations + return ArrayOperations.for_memory(self) NUMPY = ( "numpy", @@ -604,7 +600,6 @@ def _operations(self) -> ArrayOperations: (), (), FrameworkRuntime(oom_matcher=_numpy_oom), - "NUMPY_OPERATIONS", ) CUPY = ( "cupy", @@ -628,7 +623,6 @@ def _operations(self) -> ArrayOperations: oom_matcher=_cupy_oom, subprocess_environment=_nvidia_wheel_subprocess_environment, ), - "CUPY_OPERATIONS", ) TORCH = ( "torch", @@ -651,7 +645,6 @@ def _operations(self) -> ArrayOperations: dlpack_exporter=_protocol_dlpack_export, oom_matcher=_torch_oom, ), - "TORCH_OPERATIONS", ) TENSORFLOW = ( "tensorflow", @@ -670,7 +663,6 @@ def _operations(self) -> ArrayOperations: dlpack_validator=_tensorflow_dlpack, oom_matcher=_tensorflow_oom, ), - "TENSORFLOW_OPERATIONS", ) JAX = ( "jax", @@ -688,7 +680,6 @@ def _operations(self) -> ArrayOperations: dlpack_exporter=_protocol_dlpack_export, oom_matcher=_jax_oom, ), - "JAX_OPERATIONS", ) PYCLESPERANTO = ( "pyclesperanto", @@ -705,7 +696,6 @@ def _operations(self) -> ArrayOperations: move_to_active_device=_pyclesperanto_device_move, oom_matcher=_pyclesperanto_oom, ), - "PYCLESPERANTO_OPERATIONS", ) @property diff --git a/tests/test_array_operations_family.py b/tests/test_array_operations_family.py new file mode 100644 index 0000000..2450574 --- /dev/null +++ b/tests/test_array_operations_family.py @@ -0,0 +1,65 @@ +"""The ArrayOperations family: one leaf per MemoryType, found by its declaration.""" + +import ast +from pathlib import Path + +import numpy as np +import pytest + +from arraybridge import MemoryType +from arraybridge.array_operations import ArrayOperations + +SOURCE = Path(__file__).parents[1] / "src" / "arraybridge" + + +def test_every_memory_type_has_exactly_one_leaf(): + assert set(ArrayOperations.__registry__) == {member.value for member in MemoryType} + for member in MemoryType: + leaf = ArrayOperations.for_memory(member) + assert leaf.memory_type == member.value + assert ArrayOperations.for_memory(member) is leaf + assert member._operations is leaf + + +def test_leaves_are_not_named_by_strings_or_module_singletons(): + types_source = (SOURCE / "types.py").read_text() + assert "_OPERATIONS" not in types_source + tree = ast.parse((SOURCE / "array_operations.py").read_text()) + module_level_names = { + target.id + for node in tree.body + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + } + assert not {name for name in module_level_names if name.endswith("_OPERATIONS")} + + +def test_numpy_primitives(): + ops = ArrayOperations.for_memory(MemoryType.NUMPY) + stack = np.arange(24, dtype=np.uint16).reshape(2, 3, 4) + assert ops.module is np + assert ops.accepts(stack) and not ops.accepts([1, 2]) + np.testing.assert_array_equal(ops.max(stack, 0), stack[1]) + np.testing.assert_array_equal(ops.min(stack, 0), stack[0]) + np.testing.assert_array_equal(ops.mean(stack, 0), stack.mean(axis=0)) + np.testing.assert_array_equal(ops.sum(stack, (1, 2)), stack.sum(axis=(1, 2))) + assert ops.prepend_axis(stack[0]).shape == (1, 3, 4) + assert ops.astype(stack, np.float32).dtype == np.float32 + ramp = ops.linspace(0, 1, 4, endpoint=False) + np.testing.assert_array_equal(ramp, [0, 0.25, 0.5, 0.75]) + weights = ops.assign(ops.ones(5, np.float32), slice(-2, None), ramp[:2]) + np.testing.assert_array_equal(weights, [1, 1, 1, 0, 0.25]) + np.testing.assert_array_equal(ops.outer(ramp, ramp), np.outer(ramp, ramp)) + assert ops.floor(2.7) == 2 + + +def test_jax_primitives_are_functional(): + jnp = pytest.importorskip("jax.numpy") + ops = ArrayOperations.for_memory(MemoryType.JAX) + values = jnp.ones(4, dtype=jnp.float32) + updated = ops.assign(values, slice(None, 2), jnp.zeros(2, dtype=jnp.float32)) + np.testing.assert_array_equal(np.asarray(updated), [0, 0, 1, 1]) + np.testing.assert_array_equal(np.asarray(values), [1, 1, 1, 1]) + stack = jnp.arange(24, dtype=jnp.uint16).reshape(2, 3, 4) + assert ops.mean(stack, 0).dtype == jnp.float32 diff --git a/tests/test_native_geometry.py b/tests/test_native_geometry.py index 8e9a9cd..5aa0ba8 100644 --- a/tests/test_native_geometry.py +++ b/tests/test_native_geometry.py @@ -3,9 +3,11 @@ import numpy as np import pytest -from arraybridge.array_operations import PYCLESPERANTO_OPERATIONS +from arraybridge.array_operations import ArrayOperations from arraybridge.types import MemoryType +PYCLESPERANTO_OPERATIONS = ArrayOperations.for_memory(MemoryType.PYCLESPERANTO) + def test_numpy_views_and_independent_reference_allocation(): source = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] diff --git a/tests/test_stack_utils.py b/tests/test_stack_utils.py index c4f7580..863b8db 100644 --- a/tests/test_stack_utils.py +++ b/tests/test_stack_utils.py @@ -332,7 +332,10 @@ def test_tensorflow_stack_uses_immutable_framework_leaf(self, monkeypatch): assert state["device"] == 0 def test_pyclesperanto_stack_uses_binary_api_and_preserves_singleton_axis(self): - from arraybridge.array_operations import PYCLESPERANTO_OPERATIONS + from arraybridge.array_operations import ArrayOperations + from arraybridge.types import MemoryType + + operations = ArrayOperations.for_memory(MemoryType.PYCLESPERANTO) class Array: def __init__(self, values): @@ -365,8 +368,8 @@ def concatenate_along_z(cls, left, right): second = Array([[5, 6], [7, 8]]) third = Array([[9, 10], [11, 12]]) - singleton = PYCLESPERANTO_OPERATIONS.stack([first], Module) - stack = PYCLESPERANTO_OPERATIONS.stack([first, second, third], Module) + singleton = operations.stack([first], Module) + stack = operations.stack([first, second, third], Module) assert singleton.shape == (1, 2, 2) assert stack.shape == (3, 2, 2) From 6f6f907d5e012124febb7ac9063df311cc076ef2 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Sat, 10 Oct 2026 21:37:21 -0400 Subject: [PATCH 2/2] TensorFlow leaf names dtypes through its NumPy dtype; torch leaf uses its dtype-name helper The TensorFlow leaf inherited the NumPy dtype_name, so every decorated TensorFlow call failed with "Cannot interpret 'tf.uint16' as a data type"; the _tensorflow_dtype_name helper existed but nothing used it. The torch leaf restated _torch_dtype_name inline. Co-Authored-By: Claude Opus 5.5 --- src/arraybridge/array_operations.py | 6 +++++- tests/test_array_operations_family.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/arraybridge/array_operations.py b/src/arraybridge/array_operations.py index dc99c91..13772a9 100644 --- a/src/arraybridge/array_operations.py +++ b/src/arraybridge/array_operations.py @@ -390,7 +390,7 @@ def cast(data: Any, dtype: Any, module: Any) -> Any: @staticmethod def dtype_name(dtype: Any) -> str: - return str(dtype).rsplit(".", maxsplit=1)[-1] + return _torch_dtype_name(dtype) @staticmethod def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any: @@ -438,6 +438,10 @@ def outer(self, left: Any, right: Any) -> Any: def floor(self, value: float) -> Any: return self.module.math.floor(value) + @staticmethod + def dtype_name(dtype: Any) -> str: + return _tensorflow_dtype_name(dtype) + @staticmethod def to_numpy(data: Any, module: Any) -> Any: del module diff --git a/tests/test_array_operations_family.py b/tests/test_array_operations_family.py index 2450574..803cf69 100644 --- a/tests/test_array_operations_family.py +++ b/tests/test_array_operations_family.py @@ -63,3 +63,12 @@ def test_jax_primitives_are_functional(): np.testing.assert_array_equal(np.asarray(values), [1, 1, 1, 1]) stack = jnp.arange(24, dtype=jnp.uint16).reshape(2, 3, 4) assert ops.mean(stack, 0).dtype == jnp.float32 + + +def test_tensorflow_dtype_names_come_from_the_numpy_dtype(): + class TensorflowDtype: + as_numpy_dtype = np.uint16 + + ops = ArrayOperations.for_memory(MemoryType.TENSORFLOW) + assert ops.dtype_name(TensorflowDtype()) == "uint16" + assert MemoryType.TENSORFLOW.canonical_dtype_name(TensorflowDtype()) == "uint16"