From 5456ce50bf02b887160169a28fb33e25586f3731 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 15:34:35 -0400 Subject: [PATCH 1/5] No library inherits an array operation it cannot perform ArrayLibrary supplies no implementation that assumes an array API. The NumPy-style defaults move to NumpyArrayApi (NumPy, CuPy, JAX); linspace, ones, outer and floor are declared by ArrayMath, which pyclesperanto does not inherit; assignment by item is InPlaceAssignment. Co-Authored-By: Claude Opus 5.5 --- ADDING_NEW_FRAMEWORKS.md | 67 +++++++++-- README.md | 3 + pyproject.toml | 2 +- src/arraybridge/libraries.py | 191 ++++++++++++++++++++++-------- tests/test_operation_soundness.py | 171 ++++++++++++++++++++++++++ tests/test_toy_library.py | 40 +++++++ 6 files changed, 418 insertions(+), 56 deletions(-) create mode 100644 tests/test_operation_soundness.py diff --git a/ADDING_NEW_FRAMEWORKS.md b/ADDING_NEW_FRAMEWORKS.md index a2aad93..75f35d9 100644 --- a/ADDING_NEW_FRAMEWORKS.md +++ b/ADDING_NEW_FRAMEWORKS.md @@ -71,9 +71,45 @@ class FrozenLibrary(HostDevices[FrozenArray]): def _move_to_host(self, data): return data - # The inherited default calls ``self.module.stack``; this module has none. + # Array operations: nothing is inherited for them, because the + # base class assumes no array API. This array wraps NumPy values. + def accepts(self, value): + return isinstance(value, FrozenArray) + + def max(self, data, axis): + return FrozenArray(data.values.max(axis=axis)) + + def min(self, data, axis): + return FrozenArray(data.values.min(axis=axis)) + + def mean(self, data, axis): + return FrozenArray(data.values.mean(axis=axis)) + + def sum(self, data, axis): + return FrozenArray(data.values.sum(axis=axis)) + + def assign(self, data, index, values): + updated = data.values.copy() + updated[index] = values.values + return FrozenArray(updated) + def _stack(self, values): return FrozenArray(np.stack([value.values for value in values])) + + def _cast(self, data, dtype): + return FrozenArray(data.values.astype(dtype)) + + def _logical_and(self, left, right): + return FrozenArray(np.logical_and(left.values, right.values)) + + def _reshape(self, data, shape): + return FrozenArray(data.values.reshape(shape)) + + def _broadcast_to(self, data, shape): + return FrozenArray(np.broadcast_to(data.values, shape)) + + def _ones_like(self, reference, shape, dtype): + return FrozenArray(np.ones(shape, dtype=dtype)) # Lookup, by name and by array. @@ -124,7 +160,8 @@ assert first_row(frozen).shape == (3,) | `oom_messages` | Lower-case text that marks an out-of-memory error message. | `()` | | `oom_exception_paths` | Attribute paths from the module to its out-of-memory exception classes. | `()` | -Every library defines three methods asked for by `ArrayLibrary` itself: +Every library defines three methods asked for by `ArrayLibrary` itself, and +the array operations listed under "Array operations" below: - `device_of(data)`: the device holding the array. It never returns `None`. - `_to_host(data)`: the array as a NumPy array. @@ -144,7 +181,6 @@ class must define. | `DLPackImport` | can wrap another library's buffer through DLPack. | `_import_dlpack(payload)` | | `DLPackProtocolExport` | has arrays with the standard DLPack export methods on every device. | none | | `NvidiaWheelLibraries` | loads CUDA from the `nvidia-*` wheels; child processes get their library directories on the search path. | none | -| `MutableArrays` | has the NumPy array API and arrays that can be changed in place; supplies `_scale_dtype` and in-place plane normalisation. | `_rescales_into(target_dtype)` | Methods with a default worth knowing: @@ -157,11 +193,26 @@ Methods with a default worth knowing: - `exports_dlpack_from`, `exports_dlpack_data`, `export_dlpack`: override these directly when export depends on the device or the array, as TensorFlow does. -- `_stack`, `_cast`, `_reshape`, `_broadcast_to`, `_ones_like`, - `_logical_and`, `max`, `min`, `mean`, `sum`, `linspace`, `ones`, `assign`, - `outer`, `floor`, `accepts`, `nbytes_of`, `dtype_name`: the defaults call - NumPy-named functions on the library's module or methods on the array. - Override each one whose call or numeric rule differs. +- `nbytes_of`, `dtype_name`: the defaults read `shape` and `dtype` from the + array and name the dtype as NumPy does. Override where the library differs. + +## Array operations + +`ArrayLibrary` asks every library for `accepts`, `max`, `min`, `mean`, `sum`, +`assign`, `_stack`, `_cast`, `_reshape`, `_broadcast_to`, `_ones_like` and +`_logical_and`, and gives no default for any of them: the base class assumes +no array API. A library defines them or inherits them from a parent it +declares. + +| Parent | The library... | Supplies | Abstract methods | +|---|---|---|---| +| `InPlaceAssignment` | has arrays that take `array[index] = values`. | `assign` | none | +| `ArrayMath` | can make arrays without a reference array, on its default device. A library that cannot (pyclesperanto) does not inherit it and has none of these four methods. | nothing | `linspace`, `ones`, `outer`, `floor` | +| `NumpyArrayApi` | has an array module and arrays that follow the NumPy API (`array_module.max(data, axis=...)`, `data.astype(...)`, ...). It is an `ArrayMath` library. | every operation above except `assign`, and the four of `ArrayMath` | none | +| `MutableArrays` | is a `NumpyArrayApi` and `InPlaceAssignment` library; supplies `_scale_dtype` and in-place plane normalisation. | all of the above | `_rescales_into(target_dtype)` | + +Override an inherited operation only where the library's numeric rule +differs (JAX's `mean` casts to float32 first). The six built-in classes at the end of `src/arraybridge/libraries.py` are the reference for each combination of roles. diff --git a/README.md b/README.md index d149fc5..308a430 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,9 @@ that library, on the device of their input: `max`, `min`, `mean`, `sum`, `astype`, `reshape`, `broadcast_to`, `ones_like`, `ones`, `prepend_axis`, `linspace`, `assign`, `outer`, `floor`, `logical_and`, `stack_arrays`, `normalize_planes`, `scale_dtype`, `to_numpy`, `from_numpy`. +`linspace`, `ones`, `outer` and `floor` exist only on libraries that can make +an array without a reference array (`ArrayMath`: every built-in library but +pyclesperanto). ```python import numpy as np diff --git a/pyproject.toml b/pyproject.toml index d53a490..a2372d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -136,7 +136,7 @@ ignore_missing_imports = true # Optional array libraries are not importable by the checker, so what their # modules return is untyped; those values are admitted by Array registration # at runtime (see arraybridge.libraries). -warn_return_any = false +warn_return_any = true warn_unused_configs = true disallow_untyped_defs = false # Too strict for scientific code diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py index 2f16fe2..fd7bc34 100644 --- a/src/arraybridge/libraries.py +++ b/src/arraybridge/libraries.py @@ -5,7 +5,10 @@ out-of-memory recognition, child-process environment and native array operations. What several libraries share is a parent they inherit from (``HostDevices``, ``AcceleratorDevices``, ``CudaLibrary``, ``DLPackImport``, -``DLPackProtocolExport``, ``NvidiaWheelLibraries``, ``MutableArrays``). +``DLPackProtocolExport``, ``NvidiaWheelLibraries``, ``InPlaceAssignment``, +``ArrayMath``, ``NumpyArrayApi``, ``MutableArrays``). The base class holds no +default that assumes a particular library's API: a library inherits an +implementation only from a parent it declares. Adding a library is writing one class here or anywhere else: registration, lookup, the ``MemoryType`` member, the decorator and conversion follow from it. @@ -559,11 +562,10 @@ def nbytes_of(self, data: NativeT) -> int: # ----- array operations, each on the device of its input ----- + @abstractmethod def accepts(self, value: object) -> bool: """Whether ``value`` is a native array of this library.""" - return isinstance(value, self.module.ndarray) - def stack_arrays(self, arrays: Sequence[NativeT], device: Device) -> NativeT: """Stack arrays along a new leading axis on ``device``.""" @@ -625,45 +627,31 @@ def logical_and(self, left: NativeT, right: NativeT) -> NativeT: with self._array_scope(left): return self._logical_and(left, right) + @abstractmethod def max(self, data: NativeT, axis: Axis) -> NativeT: - return self.array_module.max(data, axis=axis) + """The largest value along ``axis``.""" + @abstractmethod def min(self, data: NativeT, axis: Axis) -> NativeT: - return self.array_module.min(data, axis=axis) + """The smallest value along ``axis``.""" + @abstractmethod def mean(self, data: NativeT, axis: Axis) -> NativeT: - return self.array_module.mean(data, axis=axis) + """The mean along ``axis``.""" + @abstractmethod def sum(self, data: NativeT, axis: Axis) -> NativeT: - return self.array_module.sum(data, axis=axis) + """The sum along ``axis``.""" def prepend_axis(self, data: NativeT) -> NativeT: """Add a leading axis of length one.""" return self._reshape(data, (1, *getattr(data, "shape"))) - def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NativeT: - return self.array_module.linspace(start, stop, num, endpoint=endpoint) - - def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NativeT: - """Allocate ones on the library's default device.""" - - return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) - + @abstractmethod def assign(self, data: NativeT, index: slice, values: NativeT) -> NativeT: """Return ``data`` with ``data[index]`` replaced by ``values``.""" - getattr(data, "__setitem__")(index, values) - return data - - def outer(self, left: NativeT, right: NativeT) -> NativeT: - return self.array_module.outer(left, right) - - def floor(self, value: float) -> NativeT: - """Floor a host scalar in the library's default precision.""" - - return self.array_module.floor(value) - # ----- what a leaf supplies; these run without choosing a device ----- @abstractmethod @@ -674,24 +662,29 @@ def _to_host(self, data: NativeT) -> np.ndarray: def _scale_dtype(self, result: NativeT, target_dtype: DType) -> NativeT: """The library's intensity-preserving dtype conversion.""" + @abstractmethod def _stack(self, values: Sequence[NativeT]) -> NativeT: - return self.module.stack(tuple(values), axis=0) + """Stack along a new leading axis.""" + @abstractmethod def _cast(self, data: NativeT, dtype: DType) -> NativeT: - return getattr(data, "astype")(dtype, copy=False) + """The values as ``dtype``, without rescaling.""" + @abstractmethod def _logical_and(self, left: NativeT, right: NativeT) -> NativeT: - return self.module.logical_and(left, right) + """Element-wise ``left and right``.""" + @abstractmethod def _reshape(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: - return getattr(data, "reshape")(shape) + """The same elements in ``shape``.""" + @abstractmethod def _broadcast_to(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: - return self.module.broadcast_to(data, shape) + """``data`` repeated out to ``shape``.""" + @abstractmethod def _ones_like(self, reference: NativeT, shape: tuple[int, ...], dtype: DType) -> NativeT: - del reference - return self.module.ones(shape, dtype=dtype) + """Ones of ``shape`` and ``dtype`` on the device of ``reference``.""" def _normalize_planes( self, data: NativeT, dtype: DType, scales: Sequence[float | None] @@ -972,7 +965,89 @@ def extend_child_environment(self, environment: dict[str, str]) -> dict[str, str return prepared -class MutableArrays(ArrayLibrary[NativeT]): +class InPlaceAssignment(ArrayLibrary[NativeT]): + """Libraries whose arrays take ``array[index] = values``.""" + + def assign(self, data: NativeT, index: slice, values: NativeT) -> NativeT: + getattr(data, "__setitem__")(index, values) + return data + + +class ArrayMath(ArrayLibrary[NativeT]): + """Libraries that make arrays without a reference array, on their default device. + + A library whose module has no such functions (pyclesperanto) does not + inherit this, so asking it for one fails at the attribute, not in the call. + """ + + @abstractmethod + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NativeT: + """``num`` evenly spaced values from ``start`` to ``stop``.""" + + @abstractmethod + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NativeT: + """Ones of ``shape`` and ``dtype``.""" + + @abstractmethod + def outer(self, left: NativeT, right: NativeT) -> NativeT: + """The outer product of two vectors.""" + + @abstractmethod + def floor(self, value: float) -> NativeT: + """Floor a host scalar in the library's default precision.""" + + +class NumpyArrayApi(ArrayMath[NativeT]): + """Libraries whose array module and arrays follow the NumPy API (NumPy, CuPy, JAX).""" + + def accepts(self, value: object) -> bool: + return isinstance(value, self.array_module.ndarray) + + def max(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.max(data, axis=axis) + + def min(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.min(data, axis=axis) + + def mean(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.mean(data, axis=axis) + + def sum(self, data: NativeT, axis: Axis) -> NativeT: + return self.array_module.sum(data, axis=axis) + + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NativeT: + return self.array_module.linspace(start, stop, num, endpoint=endpoint) + + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NativeT: + return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) + + def outer(self, left: NativeT, right: NativeT) -> NativeT: + return self.array_module.outer(left, right) + + def floor(self, value: float) -> NativeT: + return self.array_module.floor(value) + + def _stack(self, values: Sequence[NativeT]) -> NativeT: + return self.array_module.stack(tuple(values), axis=0) + + def _cast(self, data: NativeT, dtype: DType) -> NativeT: + return getattr(data, "astype")(dtype, copy=False) + + def _logical_and(self, left: NativeT, right: NativeT) -> NativeT: + return self.array_module.logical_and(left, right) + + def _reshape(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: + return getattr(data, "reshape")(shape) + + def _broadcast_to(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: + return self.array_module.broadcast_to(data, shape) + + def _ones_like(self, reference: NativeT, shape: tuple[int, ...], dtype: DType) -> NativeT: + del reference + return self.array_module.ones(shape, dtype=dtype) + + +class MutableArrays(NumpyArrayApi[NativeT], InPlaceAssignment[NativeT]): """Libraries with the NumPy array API whose arrays can be changed in place.""" def _normalize_planes( @@ -1109,6 +1184,8 @@ class TorchLibrary( CudaLibrary["torch.Tensor"], DLPackImport["torch.Tensor"], DLPackProtocolExport["torch.Tensor"], + ArrayMath["torch.Tensor"], + InPlaceAssignment["torch.Tensor"], ): name = "torch" display_name = "PyTorch" @@ -1189,6 +1266,12 @@ def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> to return self.module.linspace(start, stop, num, dtype=float32) return self.module.linspace(start, stop, num + 1, dtype=float32)[:-1] + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> torch.Tensor: + return self.module.ones(shape, dtype=_mapped_dtype(dtype, self.module)) + + def outer(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + return self.module.outer(left, right) + def floor(self, value: float) -> torch.Tensor: return self.module.floor(self.module.tensor(value)) @@ -1219,6 +1302,12 @@ def _cast(self, data: torch.Tensor, dtype: DType) -> torch.Tensor: native = dtype if isinstance(dtype, module.dtype) else _mapped_dtype(dtype, module) return data.to(dtype=native) + def _logical_and(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + return self.module.logical_and(left, right) + + def _reshape(self, data: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: + return data.reshape(shape) + def _broadcast_to(self, data: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: return data.expand(shape) @@ -1233,6 +1322,7 @@ class TensorflowLibrary( HostDevices["tensorflow.Tensor"], AcceleratorDevices["tensorflow.Tensor"], DLPackImport["tensorflow.Tensor"], + ArrayMath["tensorflow.Tensor"], ): name = "tensorflow" display_name = "TensorFlow" @@ -1327,6 +1417,9 @@ def linspace( return self.module.linspace(float(start), float(stop), num) return self.module.linspace(float(start), float(stop), num + 1)[:-1] + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> tensorflow.Tensor: + return self.module.ones(shape, dtype=_mapped_dtype(dtype, self.module)) + def assign( self, data: tensorflow.Tensor, index: slice, values: tensorflow.Tensor ) -> tensorflow.Tensor: @@ -1364,9 +1457,22 @@ def _cast(self, data: tensorflow.Tensor, dtype: DType) -> tensorflow.Tensor: native = dtype if isinstance(dtype, module.dtypes.DType) else _mapped_dtype(dtype, module) return module.cast(data, native) + def _stack(self, values: Sequence[tensorflow.Tensor]) -> tensorflow.Tensor: + return self.module.stack(tuple(values), axis=0) + + def _logical_and( + self, left: tensorflow.Tensor, right: tensorflow.Tensor + ) -> tensorflow.Tensor: + return self.module.logical_and(left, right) + def _reshape(self, data: tensorflow.Tensor, shape: tuple[int, ...]) -> tensorflow.Tensor: return self.module.reshape(data, shape) + def _broadcast_to( + self, data: tensorflow.Tensor, shape: tuple[int, ...] + ) -> tensorflow.Tensor: + return self.module.broadcast_to(data, shape) + def _ones_like( self, reference: tensorflow.Tensor, shape: tuple[int, ...], dtype: DType ) -> tensorflow.Tensor: @@ -1380,6 +1486,7 @@ class JaxLibrary( AcceleratorDevices["jax.Array"], DLPackImport["jax.Array"], DLPackProtocolExport["jax.Array"], + NumpyArrayApi["jax.Array"], ): name = "jax" display_name = "JAX" @@ -1454,9 +1561,6 @@ def _aligned(data: object) -> object: return copier() return data - def accepts(self, value: object) -> bool: - return isinstance(value, self.module.numpy.ndarray) - def mean(self, data: jax.Array, axis: Axis) -> jax.Array: jnp = self.module.numpy return jnp.mean(data.astype(jnp.float32), axis=axis) @@ -1492,23 +1596,16 @@ def _scale_dtype(self, result: jax.Array, target_dtype: DType) -> jax.Array: scaled = jnp.clip(scaled, *bounds) return scaled.astype(mapped) - def _stack(self, values: Sequence[jax.Array]) -> jax.Array: - return self.module.numpy.stack(tuple(values), axis=0) - def _cast(self, data: jax.Array, dtype: DType) -> jax.Array: return data.astype(_mapped_dtype(dtype, self.module.numpy)) - def _logical_and(self, left: jax.Array, right: jax.Array) -> jax.Array: - return self.module.numpy.logical_and(left, right) - - def _broadcast_to(self, data: jax.Array, shape: tuple[int, ...]) -> jax.Array: - return self.module.numpy.broadcast_to(data, shape) - def _ones_like(self, reference: jax.Array, shape: tuple[int, ...], dtype: DType) -> jax.Array: return self.module.numpy.ones(shape, dtype=dtype, device=reference.device) -class PyclesperantoLibrary(AcceleratorDevices["pyclesperanto.Array"]): +class PyclesperantoLibrary( + AcceleratorDevices["pyclesperanto.Array"], InPlaceAssignment["pyclesperanto.Array"] +): name = "pyclesperanto" oom_messages = ("cl_mem_object_allocation_failure", "cl_out_of_resources", "out of memory") diff --git a/tests/test_operation_soundness.py b/tests/test_operation_soundness.py new file mode 100644 index 0000000..754c097 --- /dev/null +++ b/tests/test_operation_soundness.py @@ -0,0 +1,171 @@ +"""No library inherits an array operation it cannot perform. + +Two checks. Where a library can be imported here, every public array +operation it has is called and compared with NumPy. For every library, +importable or not, the class structure is checked: the base class supplies no +implementation that assumes an array API, and the operations a library cannot +perform are absent from its class. +""" + +import inspect + +import numpy as np +import pytest + +from arraybridge.libraries import ( + ArrayLibrary, + ArrayMath, + InPlaceAssignment, + NumpyArrayApi, + PyclesperantoLibrary, +) + +OPERATION_OWNERS = (ArrayLibrary, ArrayMath, InPlaceAssignment) + + +def _returns_native_array(function) -> bool: + return inspect.signature(function).return_annotation == "NativeT" + + +# Every public method of the family that returns an array of the library. +ARRAY_OPERATIONS = frozenset( + name + for owner in OPERATION_OWNERS + for name, function in vars(owner).items() + if not name.startswith("_") and inspect.isfunction(function) + if _returns_native_array(function) +) +# The private steps a leaf or a role supplies for them. +ARRAY_STEPS = frozenset( + name + for name, function in vars(ArrayLibrary).items() + if name.startswith("_") and not name.startswith("__") and inspect.isfunction(function) + if _returns_native_array(function) +) + +STACK = np.arange(24, dtype=np.float32).reshape(2, 3, 4) +VECTOR = np.array([1.0, 2.0, 3.0], dtype=np.float32) +RAMP = np.array([0.0, 0.25, 0.5, 0.75], dtype=np.float32) + +# name -> (call(library, device, upload), the NumPy answer) +OPERATION_CALLS = { + "move_to_device": (lambda lib, dev, up: lib.move_to_device(up(STACK), dev), STACK), + "unrecorded_from_numpy": (lambda lib, dev, up: lib.unrecorded_from_numpy(STACK, dev), STACK), + "from_numpy": (lambda lib, dev, up: lib.from_numpy(STACK, dev), STACK), + "stack_arrays": ( + lambda lib, dev, up: lib.stack_arrays([up(STACK[0]), up(STACK[1])], dev), + STACK, + ), + "normalize_planes": ( + lambda lib, dev, up: lib.normalize_planes(up(STACK), np.float32, [2.0, None]), + np.stack([STACK[0] / 2.0, STACK[1]]), + ), + "scale_dtype": (lambda lib, dev, up: lib.scale_dtype(up(STACK), np.float32), STACK), + "astype": (lambda lib, dev, up: lib.astype(up(STACK), np.int32), STACK.astype(np.int32)), + "reshape": (lambda lib, dev, up: lib.reshape(up(STACK), (6, 4)), STACK.reshape(6, 4)), + "broadcast_to": ( + lambda lib, dev, up: lib.broadcast_to(up(STACK[:1]), (2, 3, 4)), + np.broadcast_to(STACK[:1], (2, 3, 4)), + ), + "ones_like": ( + lambda lib, dev, up: lib.ones_like(up(STACK), dtype=np.float32), + np.ones_like(STACK), + ), + "logical_and": ( + lambda lib, dev, up: lib.logical_and(up(STACK > 3), up(STACK < 10)), + np.logical_and(STACK > 3, STACK < 10), + ), + "max": (lambda lib, dev, up: lib.max(up(STACK), 0), STACK.max(axis=0)), + "min": (lambda lib, dev, up: lib.min(up(STACK), 0), STACK.min(axis=0)), + "mean": (lambda lib, dev, up: lib.mean(up(STACK), 0), STACK.mean(axis=0)), + "sum": (lambda lib, dev, up: lib.sum(up(STACK), 0), STACK.sum(axis=0)), + "prepend_axis": (lambda lib, dev, up: lib.prepend_axis(up(STACK[0])), STACK[:1]), + "assign": ( + lambda lib, dev, up: lib.assign( + up(np.ones(5, dtype=np.float32)), slice(-2, None), up(RAMP[:2]) + ), + np.array([1, 1, 1, 0, 0.25], dtype=np.float32), + ), + "linspace": (lambda lib, dev, up: lib.linspace(0, 1, 4, endpoint=False), RAMP), + "ones": (lambda lib, dev, up: lib.ones(5, np.float32), np.ones(5, dtype=np.float32)), + "outer": ( + lambda lib, dev, up: lib.outer(up(VECTOR), up(VECTOR)), + np.outer(VECTOR, VECTOR), + ), + "floor": (lambda lib, dev, up: lib.floor(2.7), np.float32(2.0)), +} + +LIBRARY_CLASSES = tuple(ArrayLibrary.__registry__.values()) + + +def test_every_array_operation_has_a_call() -> None: + assert set(OPERATION_CALLS) == ARRAY_OPERATIONS + + +@pytest.mark.parametrize("library", ArrayLibrary.all(), ids=str) +@pytest.mark.parametrize("operation", sorted(OPERATION_CALLS)) +def test_importable_library_performs_every_operation_it_has(library, operation) -> None: + if not hasattr(library, operation): + pytest.skip(f"{library.display_name} does not have {operation}") + if library.import_if_installed() is None: + pytest.skip(f"{library.display_name} is not installed") + devices = library.available_devices() + if not devices: + pytest.skip(f"{library.display_name} has no device on this machine") + device = devices[0] + call, expected = OPERATION_CALLS[operation] + + result = call(library, device, lambda values: library.from_numpy(values, device)) + + # ``floor`` of a host scalar is a scalar of the library (NumPy: ``np.float64``). + assert library.accepts(result) or np.ndim(expected) == 0 + np.testing.assert_allclose(library.to_numpy(result), expected, rtol=1e-6) + + +def test_the_base_class_assumes_no_array_api() -> None: + """An implementation in ``ArrayLibrary`` never calls the library's module.""" + + for name in ARRAY_OPERATIONS | ARRAY_STEPS | {"accepts"}: + function = vars(ArrayLibrary).get(name) + if function is None or getattr(function, "__isabstractmethod__", False): + continue + assert not {"module", "array_module"} & set(function.__code__.co_names), name + + +@pytest.mark.parametrize("library_class", LIBRARY_CLASSES, ids=lambda cls: cls.name) +def test_library_class_leaves_nothing_abstract(library_class) -> None: + assert not inspect.isabstract(library_class) + + +@pytest.mark.parametrize("library_class", LIBRARY_CLASSES, ids=lambda cls: cls.name) +def test_reference_free_operations_exist_only_on_array_math(library_class) -> None: + reference_free = { + name for name in vars(ArrayMath) if not name.startswith("_") and name in ARRAY_OPERATIONS + } + assert reference_free == {"linspace", "ones", "outer", "floor"} + for name in reference_free: + assert hasattr(library_class, name) == issubclass(library_class, ArrayMath), name + + +@pytest.mark.parametrize("library_class", LIBRARY_CLASSES, ids=lambda cls: cls.name) +def test_numpy_style_implementations_reach_only_numpy_api_libraries(library_class) -> None: + """Outside ``NumpyArrayApi`` every operation comes from the leaf or a non-NumPy role.""" + + if issubclass(library_class, NumpyArrayApi): + return + for name in ARRAY_OPERATIONS | ARRAY_STEPS | {"accepts"}: + if not hasattr(library_class, name): + continue + owner = next(base for base in library_class.__mro__ if name in vars(base)) + assert owner is not NumpyArrayApi, name + + +def test_which_libraries_follow_the_numpy_api() -> None: + assert {cls.name for cls in LIBRARY_CLASSES if issubclass(cls, NumpyArrayApi)} == { + "numpy", + "cupy", + "jax", + } + assert not issubclass(PyclesperantoLibrary, ArrayMath) + with pytest.raises(AttributeError): + ArrayLibrary.named("pyclesperanto").linspace diff --git a/tests/test_toy_library.py b/tests/test_toy_library.py index 6d35e1e..fd19e58 100644 --- a/tests/test_toy_library.py +++ b/tests/test_toy_library.py @@ -59,6 +59,46 @@ def _move_to_host(self, data): def _scale_dtype(self, result, target_dtype): return ToyArray(result.values.astype(target_dtype)) + # Array operations: nothing is inherited for them, because the + # base class assumes no array API. This array wraps NumPy values. + def accepts(self, value): + return isinstance(value, ToyArray) + + def max(self, data, axis): + return ToyArray(data.values.max(axis=axis)) + + def min(self, data, axis): + return ToyArray(data.values.min(axis=axis)) + + def mean(self, data, axis): + return ToyArray(data.values.mean(axis=axis)) + + def sum(self, data, axis): + return ToyArray(data.values.sum(axis=axis)) + + def assign(self, data, index, values): + updated = data.values.copy() + updated[index] = values.values + return ToyArray(updated) + + def _stack(self, values): + return ToyArray(np.stack([value.values for value in values])) + + def _cast(self, data, dtype): + return ToyArray(data.values.astype(dtype)) + + def _logical_and(self, left, right): + return ToyArray(np.logical_and(left.values, right.values)) + + def _reshape(self, data, shape): + return ToyArray(data.values.reshape(shape)) + + def _broadcast_to(self, data, shape): + return ToyArray(np.broadcast_to(data.values, shape)) + + def _ones_like(self, reference, shape, dtype): + return ToyArray(np.ones(shape, dtype=dtype)) + try: yield ToyLibrary finally: From 740a972f2fa868dc591ff96fd2b4c599636a7286 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 15:36:05 -0400 Subject: [PATCH 2/5] A decorated call in main memory never touches the accelerator The execution wrapper reads the device from the call's arguments. It used to ask the library for its current device whenever an accelerator existed, which for torch initialises CUDA and opens a stream for a function working on host tensors, and used the selected device's stream even when the data was on another device. Co-Authored-By: Claude Opus 5.5 --- src/arraybridge/__init__.py | 1 + src/arraybridge/libraries.py | 88 ++++++++++++++++++++-------------- src/arraybridge/stack_utils.py | 4 +- tests/fakes.py | 4 ++ tests/test_decorators.py | 88 ++++++++++++++++++++++++++++------ 5 files changed, 133 insertions(+), 52 deletions(-) diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index 99d79d9..28efb42 100644 --- a/src/arraybridge/__init__.py +++ b/src/arraybridge/__init__.py @@ -42,6 +42,7 @@ "_supports_dlpack": ".utils", } + def _library_names() -> list[str]: """The registered library names; each is also the name of its decorator.""" diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py index fd7bc34..5fdc1d8 100644 --- a/src/arraybridge/libraries.py +++ b/src/arraybridge/libraries.py @@ -431,9 +431,7 @@ def cleanup_all_loaded(cls, device: AcceleratorDevice | None = None) -> None: try: library.cleanup_loaded(device) except Exception as error: - logger.warning( - "Failed to cleanup %s GPU memory: %s", library.display_name, error - ) + logger.warning("Failed to cleanup %s GPU memory: %s", library.display_name, error) def run_with_oom_recovery( self, @@ -796,19 +794,50 @@ def _move_to_accelerator(self, data: NativeT, device: AcceleratorDevice) -> Nati def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> NativeT: """Create a native array on ``device``; called inside that device's scope.""" - def _thread_execution( - self, - ) -> tuple[AcceleratorDevice | None, AbstractContextManager[None]]: - """The device a decorated function runs on, if one is selected, and its scope.""" + def _argument_device(self, arguments: Sequence[object]) -> Device | None: + """The device of the first array of this library among a call's arguments.""" - return None, nullcontext() + if self.loaded_module() is None: + return None + for value in arguments: + if self.accepts(value): + return self.device_of(value) + return None + + def _selected_accelerator(self) -> AcceleratorDevice | None: + """The accelerator the library has selected for this thread, where it says.""" + + return None + + def _accelerator_call_scope( + self, device: AcceleratorDevice | None + ) -> AbstractContextManager[None]: + """What a call running on ``device`` is wrapped in.""" + + del device + return nullcontext() def execution_wrapper(self, func: Callable[P, R], oom_recovery: bool) -> Callable[P, R]: + """Run ``func`` on the accelerator holding its arguments; a host call is untouched. + + The device is read from the arguments. A call whose arrays are in main + memory, or that passes no array of a library that has main memory, + runs as written: nothing of the accelerator is asked or initialised. + """ + @functools.wraps(func) def on_accelerator(*args: P.args, **kwargs: P.kwargs) -> R: - if self.import_if_installed() is None or not self.accelerator_devices(): + held_on = self._argument_device((*args, *kwargs.values())) + if isinstance(held_on, AcceleratorDevice): + device: AcceleratorDevice | None = held_on + elif held_on is not None or self.host_devices(): + return func(*args, **kwargs) + elif self.import_if_installed() is None or not self.accelerator_devices(): return func(*args, **kwargs) - device, scope = self._thread_execution() + else: + # Every array of this library lives on an accelerator. + device = self._selected_accelerator() + scope = self._accelerator_call_scope(device) def execute() -> R: with scope: @@ -864,14 +893,15 @@ def thread_stream(self, device: AcceleratorDevice) -> object: ) return streams[key] - def _thread_execution( - self, - ) -> tuple[AcceleratorDevice | None, AbstractContextManager[None]]: - device = self.current_device() + def _selected_accelerator(self) -> AcceleratorDevice | None: + return self.current_device() + + def _accelerator_call_scope( + self, device: AcceleratorDevice | None + ) -> AbstractContextManager[None]: if device is None: - return None, nullcontext() - self.require_device(device) - return device, self.stream_scope(self.thread_stream(device)) + return nullcontext() + return self.stream_scope(self.thread_stream(device)) def cleanup_loaded(self, device: AcceleratorDevice | None = None) -> None: if self.loaded_module() is None: @@ -912,9 +942,7 @@ def exports_dlpack_from(self, device: Device) -> bool: return True def exports_dlpack_data(self, data: NativeT) -> bool: - return any( - callable(getattr(data, attribute, None)) for attribute in self._export_methods - ) + return any(callable(getattr(data, attribute, None)) for attribute in self._export_methods) def export_dlpack(self, data: NativeT) -> DLPackPayload | None: for attribute in self._export_methods: @@ -1381,9 +1409,7 @@ def exports_dlpack_data(self, data: tensorflow.Tensor) -> bool: def export_dlpack(self, data: tensorflow.Tensor) -> DLPackPayload | None: if not self.exports_dlpack_data(data): return None - return DLPackPayload( - source=data, capsule=self.module.experimental.dlpack.to_dlpack(data) - ) + return DLPackPayload(source=data, capsule=self.module.experimental.dlpack.to_dlpack(data)) def _import_dlpack(self, payload: DLPackPayload) -> tensorflow.Tensor: return self.module.experimental.dlpack.from_dlpack(payload.capsule) @@ -1410,9 +1436,7 @@ def mean(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: def sum(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: return self.module.reduce_sum(data, axis=axis) - def linspace( - self, start: float, stop: float, num: int, *, endpoint: bool - ) -> tensorflow.Tensor: + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> tensorflow.Tensor: if endpoint: return self.module.linspace(float(start), float(stop), num) return self.module.linspace(float(start), float(stop), num + 1)[:-1] @@ -1460,17 +1484,13 @@ def _cast(self, data: tensorflow.Tensor, dtype: DType) -> tensorflow.Tensor: def _stack(self, values: Sequence[tensorflow.Tensor]) -> tensorflow.Tensor: return self.module.stack(tuple(values), axis=0) - def _logical_and( - self, left: tensorflow.Tensor, right: tensorflow.Tensor - ) -> tensorflow.Tensor: + def _logical_and(self, left: tensorflow.Tensor, right: tensorflow.Tensor) -> tensorflow.Tensor: return self.module.logical_and(left, right) def _reshape(self, data: tensorflow.Tensor, shape: tuple[int, ...]) -> tensorflow.Tensor: return self.module.reshape(data, shape) - def _broadcast_to( - self, data: tensorflow.Tensor, shape: tuple[int, ...] - ) -> tensorflow.Tensor: + def _broadcast_to(self, data: tensorflow.Tensor, shape: tuple[int, ...]) -> tensorflow.Tensor: return self.module.broadcast_to(data, shape) def _ones_like( @@ -1710,9 +1730,7 @@ def prepend_axis(self, data: pyclesperanto.Array) -> pyclesperanto.Array: def _to_host(self, data: pyclesperanto.Array) -> np.ndarray: return self.module.pull(data) - def _scale_dtype( - self, result: pyclesperanto.Array, target_dtype: DType - ) -> pyclesperanto.Array: + def _scale_dtype(self, result: pyclesperanto.Array, target_dtype: DType) -> pyclesperanto.Array: module, np = self.module, _numpy() target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer) if not (np.issubdtype(result.dtype, np.floating) and target_is_int): diff --git a/src/arraybridge/stack_utils.py b/src/arraybridge/stack_utils.py index e810f96..b799758 100644 --- a/src/arraybridge/stack_utils.py +++ b/src/arraybridge/stack_utils.py @@ -154,9 +154,7 @@ def unstack_slices( array = library.move_to_device(array, device) # Extract slices along axis 0 (already in the target memory type) - slices = [ - getattr(array, "__getitem__")(index) for index in range(getattr(array, "shape")[0]) - ] + slices = [getattr(array, "__getitem__")(index) for index in range(getattr(array, "shape")[0])] # Validate that all extracted slices are 2D if requested if validate_slices: diff --git a/tests/fakes.py b/tests/fakes.py index e29e0ab..22ee7c4 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -104,6 +104,8 @@ def __dlpack_device__(self): class FakeCupy: + ndarray = FakeCupyArray + def __init__(self, device_count=2): self.events = [] self.current = 0 @@ -162,6 +164,8 @@ def nelement(self): class FakeTorch: """torch with main memory and ``device_count`` CUDA devices (none by default).""" + Tensor = FakeTorchTensor + def __init__(self, device_count=0): self.events = [] self.current = 0 diff --git a/tests/test_decorators.py b/tests/test_decorators.py index 7aad9ec..d172a98 100644 --- a/tests/test_decorators.py +++ b/tests/test_decorators.py @@ -23,7 +23,7 @@ from arraybridge.devices import AcceleratorDevice from arraybridge.libraries import ArrayLibrary from arraybridge.types import MemoryType -from tests.fakes import FakeCupy, FakeTorch +from tests.fakes import FakeCupy, FakeCupyArray, FakeTorch, FakeTorchTensor class TestDtypeConversion: @@ -63,31 +63,93 @@ class TestExecutionOnAccelerators: """What a library's devices need around a decorated call.""" @pytest.mark.parametrize( - ("name", "fake"), [("torch", FakeTorch), ("cupy", FakeCupy)], ids=["torch", "cupy"] + ("name", "fake", "on_device"), + [ + ("torch", FakeTorch, lambda module, index: FakeTorchTensor([1], cuda_index=index)), + ("cupy", FakeCupy, lambda module, index: FakeCupyArray([1], index, module)), + ], + ids=["torch", "cupy"], ) - def test_cuda_call_runs_on_this_threads_stream(self, name, fake, monkeypatch): - module = fake(device_count=1) + def test_cuda_call_runs_on_the_stream_of_the_device_holding_its_arguments( + self, name, fake, on_device, monkeypatch + ): + module = fake(device_count=2) monkeypatch.setitem(sys.modules, name, module) + # The data is on device 1 while the process has device 0 selected. + payload, again = on_device(module, 1), on_device(module, 1) @getattr(decorators, name) def identity(value): module.events.append(("call", value)) return value - assert identity("payload") == "payload" - assert identity("again") == "again" - stream = ArrayLibrary.named(name).thread_stream(AcceleratorDevice(0)) + assert identity(payload) is payload + assert identity(again) is again + stream = ArrayLibrary.named(name).thread_stream(AcceleratorDevice(1)) assert module.events == [ - ("enter", 0), - ("exit", 0), + ("enter", 1), + ("exit", 1), ("stream enter", stream), - ("call", "payload"), + ("call", payload), ("stream exit", stream), ("stream enter", stream), - ("call", "again"), + ("call", again), + ("stream exit", stream), + ] + + def test_cupy_call_without_an_array_runs_on_the_selected_device(self, monkeypatch): + # CuPy has no main memory, so a call always runs on an accelerator. + module = FakeCupy(device_count=2) + module.current = 1 + monkeypatch.setitem(sys.modules, "cupy", module) + + @decorators.cupy + def identity(value): + module.events.append(("call", value)) + return value + + assert identity("payload") == "payload" + stream = ArrayLibrary.named("cupy").thread_stream(AcceleratorDevice(1)) + assert module.events[-3:] == [ + ("stream enter", stream), + ("call", "payload"), ("stream exit", stream), ] + @pytest.mark.parametrize( + "arguments", + [ + lambda: (FakeTorchTensor([1, 2]),), + lambda: (np.array([1, 2]),), + lambda: ("no array at all",), + ], + ids=["host tensor", "numpy array", "no array"], + ) + def test_torch_call_in_main_memory_never_touches_cuda(self, arguments, monkeypatch): + class TorchWithUntouchableCuda: + Tensor = FakeTorchTensor + + @property + def cuda(self): + raise AssertionError("a call in main memory asked torch.cuda for something") + + monkeypatch.setitem(sys.modules, "torch", TorchWithUntouchableCuda()) + calls = [] + + @decorators.torch + def fails_once(value): + calls.append(value) + if len(calls) == 1: + raise MemoryError("out of memory") + return value + + # Not even an out-of-memory error makes a host call clear accelerator caches. + with pytest.raises(MemoryError): + fails_once(*arguments()) + (value,) = arguments() + assert fails_once(value) is value + assert len(calls) == 2 + def test_torch_without_cuda_calls_the_function_directly(self, monkeypatch): module = FakeTorch() monkeypatch.setitem(sys.modules, "torch", module) @@ -179,9 +241,7 @@ def declared(image, scale=3): assert restored.__name__ == declared.__name__ assert restored.__wrapped__.__name__ == declared.__wrapped__.__name__ assert restored.execution_memory_type == name - np.testing.assert_array_equal( - restored(np.array([1, 2], dtype=np.uint16), scale=4), [4, 8] - ) + np.testing.assert_array_equal(restored(np.array([1, 2], dtype=np.uint16), scale=4), [4, 8]) assert ArrayLibrary.named("cupy").thread_stream(AcceleratorDevice(0)) is stream From 6820331f23c43e72159ef74f32ebf0fff6a0c300 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 15:36:13 -0400 Subject: [PATCH 3/5] CONTRIBUTING: drop the compatibility rule Co-Authored-By: Claude Opus 5.5 --- CONTRIBUTING.md | 1 - 1 file changed, 1 deletion(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 80455c2..75611b4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -403,7 +403,6 @@ GitHub Actions will automatically: 2. **Start small**: Begin with MVP, iterate based on feedback 3. **Document**: Add docstrings, update README/docs 4. **Test thoroughly**: Include unit, integration, and edge case tests -5. **Maintain compatibility**: Avoid breaking changes when possible ### Adding Library Support From e73048127a4ef90b574a1ed0960bb4db8d7e4658 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 15:57:52 -0400 Subject: [PATCH 4/5] mypy reports every returned Any; each library's module has a declared type warn_return_any is on and ignore_missing_imports is off. The numpy.* follow_imports = "skip" override is gone. Every library the package imports is declared in stubs/, which lists what arraybridge calls and comes before an installed copy on mypy_path, so the result does not depend on what is installed. CI checks `mypy src/ stubs/`. - A leaf states `module = ` under TYPE_CHECKING, so its calls are checked against the declaration of that one library. - NumpyArrayApi and MutableArrays name their module by the Protocols NumpyFunctions and MutableNumpyFunctions, and their arrays by NumpyArray and MutableNumpyArray; each leaf returns its module from `array_module` and the checker compares the two. - Values with no static type are checked where they enter: _mapped_dtype (the dtype class asked for), MemoryContractAttribute.read (a name or None), DeviceMove.convert (one library), DLPackExchange.convert (the target imports DLPack), TorchLibrary.stream_scope (a torch stream). - AcceleratorDevices asks each leaf for _device_if_accepted; CudaLibrary for create_stream; DLPackImport.from_dlpack returns NotImplemented when the leaf's _reads_dlpack says no, so _import_dlpack returns an array. - The MemoryType metaclass overrides carry EnumMeta's signatures. The __call__ override is removed: Enum already asks _missing_ for a value it does not know. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 4 +- ADDING_NEW_FRAMEWORKS.md | 40 +- CHANGELOG.md | 24 + CONTRIBUTING.md | 6 +- TESTING.md | 2 +- docs/source/ci-cd.rst | 2 +- docs/source/contributing.rst | 2 +- pyproject.toml | 17 +- src/arraybridge/conversion.py | 72 ++- src/arraybridge/decorators.py | 4 +- src/arraybridge/libraries.py | 534 ++++++++++++++------- src/arraybridge/types.py | 131 +++-- stubs/cupy/__init__.pyi | 52 ++ stubs/cupy/cuda/__init__.pyi | 29 ++ stubs/cupy/cuda/runtime.pyi | 5 + stubs/jax/__init__.pyi | 41 ++ stubs/jax/dlpack.pyi | 5 + stubs/jax/numpy.pyi | 33 ++ stubs/metaclass_registry/__init__.pyi | 12 + stubs/numpy/__init__.pyi | 61 +++ stubs/pyclesperanto/__init__.pyi | 32 ++ stubs/python_introspect/__init__.pyi | 15 + stubs/tensorflow/__init__.pyi | 54 +++ stubs/tensorflow/config.pyi | 5 + stubs/tensorflow/dtypes.pyi | 4 + stubs/tensorflow/experimental/__init__.pyi | 3 + stubs/tensorflow/experimental/dlpack.pyi | 6 + stubs/tensorflow/math.pyi | 5 + stubs/torch/__init__.pyi | 58 +++ stubs/torch/cuda.pyi | 33 ++ tests/fakes.py | 10 +- tests/test_checked_values.py | 145 ++++++ 32 files changed, 1197 insertions(+), 249 deletions(-) create mode 100644 stubs/cupy/__init__.pyi create mode 100644 stubs/cupy/cuda/__init__.pyi create mode 100644 stubs/cupy/cuda/runtime.pyi create mode 100644 stubs/jax/__init__.pyi create mode 100644 stubs/jax/dlpack.pyi create mode 100644 stubs/jax/numpy.pyi create mode 100644 stubs/metaclass_registry/__init__.pyi create mode 100644 stubs/numpy/__init__.pyi create mode 100644 stubs/pyclesperanto/__init__.pyi create mode 100644 stubs/python_introspect/__init__.pyi create mode 100644 stubs/tensorflow/__init__.pyi create mode 100644 stubs/tensorflow/config.pyi create mode 100644 stubs/tensorflow/dtypes.pyi create mode 100644 stubs/tensorflow/experimental/__init__.pyi create mode 100644 stubs/tensorflow/experimental/dlpack.pyi create mode 100644 stubs/tensorflow/math.pyi create mode 100644 stubs/torch/__init__.pyi create mode 100644 stubs/torch/cuda.pyi create mode 100644 tests/test_checked_values.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9f0ea64..4131829 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -152,6 +152,8 @@ jobs: run: | black --check src/ + # stubs/ declares what arraybridge calls of each library it imports, so + # nothing but mypy needs installing; naming it checks the declarations too. - name: Run mypy (type checking) run: | - mypy src/ + mypy src/ stubs/ diff --git a/ADDING_NEW_FRAMEWORKS.md b/ADDING_NEW_FRAMEWORKS.md index 75f35d9..3911db1 100644 --- a/ADDING_NEW_FRAMEWORKS.md +++ b/ADDING_NEW_FRAMEWORKS.md @@ -176,8 +176,8 @@ class must define. | Parent | The library... | Abstract methods | |---|---|---| | `HostDevices` | keeps arrays in main memory. | `_wrap_host_array(data)`, `_move_to_host(data)` | -| `AcceleratorDevices` | can keep arrays on numbered accelerators. | `_accelerator_count()`, `_accelerator_scope(device)`, `_move_to_accelerator(data, device)`, `_upload(data, device)` | -| `CudaLibrary` | is an `AcceleratorDevices` library with its own CUDA streams and a memory cache to release. | those of `AcceleratorDevices`, plus `stream_scope(stream)`, `current_device()`, `_free_cached_memory()` | +| `AcceleratorDevices` | can keep arrays on numbered accelerators. | `_accelerator_count()`, `_accelerator_scope(device)`, `_move_to_accelerator(data, device)`, `_upload(data, device)`, `_device_if_accepted(value)` | +| `CudaLibrary` | is an `AcceleratorDevices` library with its own CUDA streams and a memory cache to release. | those of `AcceleratorDevices`, plus `create_stream()`, `stream_scope(stream)`, `current_device()`, `_free_cached_memory()` | | `DLPackImport` | can wrap another library's buffer through DLPack. | `_import_dlpack(payload)` | | `DLPackProtocolExport` | has arrays with the standard DLPack export methods on every device. | none | | `NvidiaWheelLibraries` | loads CUDA from the `nvidia-*` wheels; child processes get their library directories on the search path. | none | @@ -190,6 +190,9 @@ Methods with a default worth knowing: whole process. The default does nothing. - `cleanup_loaded(device=None)`: release cached device memory. `CudaLibrary` implements it; elsewhere the default does nothing. +- `DLPackImport._reads_dlpack(payload)`: whether the importer can read this + payload. The default says yes; where it says no, `from_dlpack` returns + `NotImplemented` and `_import_dlpack` is not called. - `exports_dlpack_from`, `exports_dlpack_data`, `export_dlpack`: override these directly when export depends on the device or the array, as TensorFlow does. @@ -208,8 +211,8 @@ declares. |---|---|---|---| | `InPlaceAssignment` | has arrays that take `array[index] = values`. | `assign` | none | | `ArrayMath` | can make arrays without a reference array, on its default device. A library that cannot (pyclesperanto) does not inherit it and has none of these four methods. | nothing | `linspace`, `ones`, `outer`, `floor` | -| `NumpyArrayApi` | has an array module and arrays that follow the NumPy API (`array_module.max(data, axis=...)`, `data.astype(...)`, ...). It is an `ArrayMath` library. | every operation above except `assign`, and the four of `ArrayMath` | none | -| `MutableArrays` | is a `NumpyArrayApi` and `InPlaceAssignment` library; supplies `_scale_dtype` and in-place plane normalisation. | all of the above | `_rescales_into(target_dtype)` | +| `NumpyArrayApi` | has an array module and arrays that follow the NumPy API (`array_module.max(data, axis=...)`, `data.astype(...)`, ...). It is an `ArrayMath` library. | every operation above except `assign`, and the four of `ArrayMath` | `array_module` | +| `MutableArrays` | is a `NumpyArrayApi` and `InPlaceAssignment` library; supplies `_scale_dtype` and in-place plane normalisation. | all of the above | `array_module`, `_rescales_into(target_dtype)` | Override an inherited operation only where the library's numeric rule differs (JAX's `mean` casts to float32 first). @@ -217,6 +220,33 @@ differs (JAX's `mean` casts to float32 first). The six built-in classes at the end of `src/arraybridge/libraries.py` are the reference for each combination of roles. +## Typing + +The checker reads a library from `stubs//`, never from an installed +copy, so `mypy` gives the same result on every machine. For a new library: + +- declare in `stubs//` exactly the functions, classes and attributes + the class calls, with the library's own names and types; +- state the module in the class, so its calls are checked against that + declaration: + + ```python + if TYPE_CHECKING: + import mylib + + class MylibLibrary(HostDevices["mylib.Array"]): + name = "mylib" + + if TYPE_CHECKING: + module = mylib + ``` + +- call `self.module.(...)` directly: a local `module = self.module` + is a plain `ModuleType` to the checker and loses the declaration; +- a `NumpyArrayApi` library returns its module of NumPy-named functions from + `array_module`; the checker compares that module with the `NumpyFunctions` + Protocol (`MutableNumpyFunctions` for `MutableArrays`). + ## Tests to add - detection by array and by module name, and behaviour when the package is @@ -232,7 +262,7 @@ Run: ```bash ruff check src tests black --check src tests -mypy src --ignore-missing-imports +mypy src stubs pytest python -m sphinx -E -W --keep-going -b html docs/source /tmp/arraybridge-docs ``` diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c92d91..3544e41 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -120,3 +120,27 @@ cannot wrap one; the result (a host copy) is unchanged and the warning is gone. nominal type of an array of any supported library; `library.array_class` is an array of that library. Native array classes are registered with them when a library first meets one. + +`mypy` runs with `warn_return_any` and without `ignore_missing_imports`. Every +library the package imports is declared in the repository's `stubs/` +directory, which lists what arraybridge calls and comes before any installed +copy, so the result is the same on every machine. A library class states that +its `module` is that library; `NumpyArrayApi` names its module by the +`NumpyFunctions` Protocol. For a class written outside the package this means: +`NumpyArrayApi` and `MutableArrays` leave `array_module` abstract, +`AcceleratorDevices` asks for `_device_if_accepted(value)`, `CudaLibrary` for +`create_stream()`, and a `DLPackImport` library that cannot read a payload says +so in `_reads_dlpack(payload)` instead of returning `NotImplemented` from +`_import_dlpack`. + +Annotations that now say what was already returned: a device or stream scope +is an `AbstractContextManager[object]` (CuPy's device and stream are their own +context managers), `from_dlpack` returns the array or `NotImplementedType`, +and `ArrayLibrary.array_module` is an `object` in the base class. + +A value of the wrong kind is now refused with `TypeError` in four places: +`MemoryContractAttribute.read` when the stored value is neither a string nor +`None`, `TorchLibrary.stream_scope` for an object that is neither a +`torch.cuda.Stream` nor `None`, a dtype lookup in torch or TensorFlow whose +attribute is not a dtype of that library, and `DeviceMove.convert` called with +two different libraries. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 75611b4..5315fcb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -250,7 +250,7 @@ ruff check src/ tests/ --fix ```bash # Run type checker -mypy src/ --ignore-missing-imports +mypy src stubs # Check specific file mypy src/arraybridge/converters.py @@ -274,7 +274,7 @@ pytest pytest --cov=arraybridge --cov-report=term # 5. Type check (optional but recommended) -mypy src/ --ignore-missing-imports +mypy src stubs ``` ## Submitting Changes @@ -355,7 +355,7 @@ Add release notes to `CHANGELOG.md` (if exists) or create GitHub release notes. ```bash ruff check src tests scripts black --check src scripts -mypy src --ignore-missing-imports +mypy src stubs pytest python -m sphinx -E -W --keep-going -b html docs/source /tmp/arraybridge-docs git commit diff --git a/TESTING.md b/TESTING.md index 752c3b6..318784b 100644 --- a/TESTING.md +++ b/TESTING.md @@ -9,7 +9,7 @@ are the authority for the current CI matrix. python -m pip install -e ".[dev,docs]" ruff check src tests black --check src tests -mypy src --ignore-missing-imports +mypy src stubs pytest python -m sphinx -E -W --keep-going -b html docs/source /tmp/arraybridge-docs ``` diff --git a/docs/source/ci-cd.rst b/docs/source/ci-cd.rst index e24d080..422aeaa 100644 --- a/docs/source/ci-cd.rst +++ b/docs/source/ci-cd.rst @@ -12,7 +12,7 @@ Before opening a change, run: python -m pytest ruff check src tests black --check src tests - mypy src --ignore-missing-imports + mypy src stubs Documentation changes should also build with warnings as errors: diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst index ce61bb6..d659e04 100644 --- a/docs/source/contributing.rst +++ b/docs/source/contributing.rst @@ -9,7 +9,7 @@ by the repository: python -m pip install -e ".[dev,docs]" ruff check src tests black --check src tests - mypy src --ignore-missing-imports + mypy src stubs python -m pytest python -m sphinx -E -W --keep-going -b html docs/source docs/source/_build/html diff --git a/pyproject.toml b/pyproject.toml index a2372d1..a19aeec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,17 +129,18 @@ target-version = "py310" [tool.ruff.lint] select = ["E", "F", "I", "N", "W", "UP"] +[tool.ruff.lint.per-file-ignores] +# The declarations in stubs/ carry the libraries' own names. +"stubs/**" = ["N801", "N802"] + [tool.mypy] python_version = "3.10" follow_imports_for_stubs = true -ignore_missing_imports = true -# Optional array libraries are not importable by the checker, so what their -# modules return is untyped; those values are admitted by Array registration -# at runtime (see arraybridge.libraries). +ignore_missing_imports = false +# The optional array libraries are read from the declarations in stubs/, which +# list what arraybridge calls and come before any installed copy, so the result +# does not depend on what is installed where mypy runs. +mypy_path = "$MYPY_CONFIG_FILE_DIR/stubs" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = false # Too strict for scientific code - -[[tool.mypy.overrides]] -module = "numpy.*" -follow_imports = "skip" diff --git a/src/arraybridge/conversion.py b/src/arraybridge/conversion.py index 442758c..5f4bc75 100644 --- a/src/arraybridge/conversion.py +++ b/src/arraybridge/conversion.py @@ -14,14 +14,15 @@ from collections.abc import Iterator, Mapping from contextlib import contextmanager from dataclasses import asdict, dataclass, field -from typing import TYPE_CHECKING, ClassVar, TypeVar +from types import NotImplementedType +from typing import TYPE_CHECKING, ClassVar, TypeGuard, TypeVar from metaclass_registry import AutoRegisterMeta from arraybridge.devices import Device if TYPE_CHECKING: - from arraybridge.libraries import ArrayLibrary + from arraybridge.libraries import ArrayLibrary, DLPackImport logger = logging.getLogger(__name__) @@ -29,6 +30,19 @@ T = TypeVar("T") +def _is_array_of_both(data: S, source: ArrayLibrary[S], target: ArrayLibrary[T]) -> TypeGuard[T]: + """Whether ``data`` of ``source`` is an array of ``target`` too: they are one library.""" + + del data + return source is target + + +def _imports_dlpack(library: ArrayLibrary[T]) -> TypeGuard[DLPackImport[T]]: + from arraybridge.libraries import DLPackImport + + return isinstance(library, DLPackImport) + + class RouteDeclinedError(Exception): """A route that applied by declaration could not carry this particular array.""" @@ -134,11 +148,25 @@ class DeviceMove(ConversionRoute): mechanism = "device move" @classmethod - def applies(cls, source, source_device, target, target_device) -> bool: + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> bool: return source is target @classmethod - def convert(cls, data, source, target, target_device): + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: + if not _is_array_of_both(data, source, target): + raise TypeError(f"{source} and {target} are different libraries") return target.move_to_device(data, target_device) @@ -148,20 +176,34 @@ class DLPackExchange(ConversionRoute): mechanism = "dlpack" @classmethod - def applies(cls, source, source_device, target, target_device) -> bool: + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> bool: return target.imports_dlpack_from(source_device) and source.exports_dlpack_from( source_device ) @classmethod - def convert(cls, data, source, target, target_device): + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: try: + if not _imports_dlpack(target): + raise TypeError(f"{target} does not import DLPack") payload = source.export_dlpack(data) if payload is None: raise RouteDeclinedError("the source array has no DLPack export") with target.device_scope(target_device): result = target.from_dlpack(payload) - if result is NotImplemented: + if isinstance(result, NotImplementedType): raise RouteDeclinedError("the target cannot import this capsule") return target.move_to_device(result, target_device) except RouteDeclinedError: @@ -182,11 +224,23 @@ class HostCopy(ConversionRoute): mechanism = "host copy" @classmethod - def applies(cls, source, source_device, target, target_device) -> bool: + def applies( + cls, + source: ArrayLibrary[S], + source_device: Device, + target: ArrayLibrary[T], + target_device: Device, + ) -> bool: return True @classmethod - def convert(cls, data, source, target, target_device): + def convert( + cls, + data: S, + source: ArrayLibrary[S], + target: ArrayLibrary[T], + target_device: Device, + ) -> T: return target.unrecorded_from_numpy(source.unrecorded_to_numpy(data), target_device) diff --git a/src/arraybridge/decorators.py b/src/arraybridge/decorators.py index 8aa1852..e2d04e9 100644 --- a/src/arraybridge/decorators.py +++ b/src/arraybridge/decorators.py @@ -17,7 +17,7 @@ import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from enum import Enum from typing import ClassVar, ParamSpec, TypeVar @@ -168,6 +168,8 @@ class DtypeConversionRunner( """Registered dtype conversion behavior selected by DtypeConversion.""" __registry_family__ = RegistryFamily(RegistryKeyAttribute.STRATEGY_LABEL) + # Installed by the metaclass: each runner class under its conversion's value. + __registry__: ClassVar["Mapping[str, type[DtypeConversionRunner]]"] dtype_conversion: ClassVar[DtypeConversion | None] = None diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py index 5fdc1d8..432edff 100644 --- a/src/arraybridge/libraries.py +++ b/src/arraybridge/libraries.py @@ -18,10 +18,16 @@ array classes are admitted with ``ABC.register`` when the library first meets one, so no optional library is imported to define them. Inside a leaf the native class is the type parameter, named through a ``TYPE_CHECKING`` import -because the optional package cannot be imported at runtime here; a checker -without that package installed sees it as unknown. A library's module is a -``types.ModuleType``: what it returns has no static type, so those values are -checked by ``Array`` registration at runtime, not by the checker. +because the optional package cannot be imported at runtime here. + +The checker reads each optional library from the declarations in the +repository's ``stubs`` directory, which list what this module calls and take +precedence over an installed copy, so the result of checking does not depend +on what is installed. A leaf states that its ``module`` is that library, so +its calls are checked against those declarations. A role shared by libraries +with the same function names (``NumpyArrayApi``) names its module by a +``Protocol`` of the functions it calls, and each leaf's module is checked +against that Protocol. Importing this module imports no array library, NumPy included. """ @@ -43,8 +49,8 @@ from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from pathlib import Path -from types import ModuleType -from typing import TYPE_CHECKING, ClassVar, Generic, ParamSpec, TypeVar +from types import ModuleType, NotImplementedType +from typing import TYPE_CHECKING, ClassVar, Generic, ParamSpec, Protocol, TypeVar from metaclass_registry import AutoRegisterMeta @@ -54,13 +60,14 @@ from arraybridge.devices import HOST, AcceleratorDevice, Device, HostDevice, NamedDevice if TYPE_CHECKING: - # Optional packages: importable by a type checker that has them, never at runtime here. + # Optional packages, as declared in ``stubs``: never imported at runtime here. import cupy import jax import numpy as np import pyclesperanto import tensorflow import torch + from typing_extensions import Self logger = logging.getLogger(__name__) @@ -80,14 +87,95 @@ } -def _numpy() -> ModuleType: - """NumPy, imported when an operation first needs it.""" +NativeT = TypeVar("NativeT") +TargetT = TypeVar("TargetT") +DTypeT = TypeVar("DTypeT") - return importlib.import_module("numpy") +class Intensities(Protocol): + """The arithmetic that rescaling an array's intensities uses.""" -NativeT = TypeVar("NativeT") -TargetT = TypeVar("TargetT") + def __sub__(self, other: Self | float, /) -> Self: ... + + def __mul__(self, other: float, /) -> Self: ... + + def __truediv__(self, other: Self | float, /) -> Self: ... + + +IntensitiesT = TypeVar("IntensitiesT", bound=Intensities) + + +class NumpyArray(Protocol): + """An array with NumPy's own methods, as ``NumpyArrayApi`` calls them.""" + + def astype(self, dtype: DType, /, *, copy: bool = ...) -> Self: ... + + def reshape(self, shape: tuple[int, ...], /) -> Self: ... + + +class MutableNumpyArray(NumpyArray, Intensities, Protocol): + """A NumPy-style array that ``MutableArrays`` rescales and divides in place.""" + + @property + def dtype(self) -> object: ... + + def __getitem__(self, index: int, /) -> Self: ... + + def __setitem__(self, index: int, value: Self, /) -> None: ... + + def __itruediv__(self, other: Self | float, /) -> Self: ... + + def __le__(self, other: Self, /) -> object: ... + + +NumpyNativeT = TypeVar("NumpyNativeT", bound=NumpyArray) +MutableNativeT = TypeVar("MutableNativeT", bound=MutableNumpyArray) + + +class NumpyFunctions(Protocol[NumpyNativeT]): + """A module with NumPy's function names, as ``NumpyArrayApi`` calls them.""" + + @property + def ndarray(self) -> type: ... + + def max(self, a: NumpyNativeT, /, *, axis: Axis | None = ...) -> NumpyNativeT: ... + + def min(self, a: NumpyNativeT, /, *, axis: Axis | None = ...) -> NumpyNativeT: ... + + def mean(self, a: NumpyNativeT, /, *, axis: Axis | None = ...) -> NumpyNativeT: ... + + def sum(self, a: NumpyNativeT, /, *, axis: Axis | None = ...) -> NumpyNativeT: ... + + def linspace( + self, start: float, stop: float, num: int, /, *, endpoint: bool = ... + ) -> NumpyNativeT: ... + + def ones(self, shape: int | tuple[int, ...], /, *, dtype: DType = ...) -> NumpyNativeT: ... + + def outer(self, a: NumpyNativeT, b: NumpyNativeT, /) -> NumpyNativeT: ... + + def floor(self, x: float, /) -> NumpyNativeT: ... + + def stack(self, arrays: Sequence[NumpyNativeT], /, *, axis: int = ...) -> NumpyNativeT: ... + + def logical_and(self, x1: NumpyNativeT, x2: NumpyNativeT, /) -> NumpyNativeT: ... + + def broadcast_to(self, array: NumpyNativeT, shape: tuple[int, ...], /) -> NumpyNativeT: ... + + +class MutableNumpyFunctions(NumpyFunctions[MutableNativeT], Protocol[MutableNativeT]): + """The further NumPy function names that ``MutableArrays`` calls.""" + + @property + def floating(self) -> object: ... + + def issubdtype(self, arg1: object, arg2: object, /) -> bool: ... + + def array( + self, a: MutableNativeT, /, *, dtype: DType = ..., copy: bool = ... + ) -> MutableNativeT: ... + + def clip(self, a: MutableNativeT, a_min: float, a_max: float, /) -> MutableNativeT: ... @dataclass(frozen=True, slots=True) @@ -105,7 +193,12 @@ def _dtype_name(dtype: DType) -> str: return str(getattr(dtype, "__name__", str(dtype).rsplit(".", maxsplit=1)[-1])) -def _scaled_values(result, result_min, result_max, target_dtype: DType): +def _scaled_values( + result: IntensitiesT, + result_min: IntensitiesT, + result_max: IntensitiesT, + target_dtype: DType, +) -> IntensitiesT: normalized = (result - result_min) / (result_max - result_min) range_info = _SCALING_RANGES.get(_dtype_name(target_dtype)) if range_info is None: @@ -126,14 +219,21 @@ def _clamp_bounds(target_dtype: DType) -> tuple[float, float] | None: return 0, range_info -def _mapped_dtype(target_dtype: DType, module: ModuleType): +def _mapped_dtype(target_dtype: DType, module: object, dtype_class: type[DTypeT]) -> DTypeT: + """The ``dtype_class`` object that ``module`` keeps under the NumPy name of ``target_dtype``.""" + + import numpy as np + try: - dtype_name = _numpy().dtype(target_dtype).name + dtype_name = np.dtype(target_dtype).name except TypeError as error: raise TypeError(f"Unsupported target dtype {target_dtype!r}") from error - mapped = getattr(module, dtype_name, None) + module_name = getattr(module, "__name__", module) + mapped: object = getattr(module, dtype_name, None) if mapped is None: - raise TypeError(f"{module.__name__} does not expose dtype {dtype_name}") + raise TypeError(f"{module_name} does not expose dtype {dtype_name}") + if not isinstance(mapped, dtype_class): + raise TypeError(f"{module_name}.{dtype_name} is not a {dtype_class.__name__}") return mapped @@ -299,7 +399,7 @@ def module(self) -> ModuleType: return self.loaded_module() or self.import_module() @property - def array_module(self) -> ModuleType: + def array_module(self) -> object: """The module holding the library's array functions.""" return self.module @@ -349,7 +449,7 @@ def require_device(self, device: Device) -> None: f"available devices are {tuple(each.label for each in available)}" ) - def device_scope(self, device: Device) -> AbstractContextManager[None]: + def device_scope(self, device: Device) -> AbstractContextManager[object]: """A block in which new arrays of this library are created on ``device``.""" self.require_device(device) @@ -370,7 +470,7 @@ def move_to_device(self, data: NativeT, device: Device) -> NativeT: with self._scope(device): return self._move(data, device) - def _array_scope(self, reference: NativeT) -> AbstractContextManager[None]: + def _array_scope(self, reference: NativeT) -> AbstractContextManager[object]: """The scope of the accelerator holding ``reference``; host data needs none.""" device = self.device_of(reference) @@ -380,7 +480,7 @@ def _array_scope(self, reference: NativeT) -> AbstractContextManager[None]: # Each device parent answers for its own kind of device and passes the rest on. - def _scope(self, device: Device) -> AbstractContextManager[None]: + def _scope(self, device: Device) -> AbstractContextManager[object]: raise ValueError(f"{self.display_name} has no {device.label} device") def _activate(self, device: Device) -> None: @@ -553,7 +653,8 @@ def convert_to(self, data: NativeT, target: ArrayLibrary[TargetT], device: Devic def nbytes_of(self, data: NativeT) -> int: """Size of the array's elements in bytes.""" - np = _numpy() + import numpy as np + return int(np.prod(tuple(getattr(data, "shape")), dtype=np.int64)) * int( np.dtype(getattr(data, "dtype")).itemsize ) @@ -591,7 +692,9 @@ def scale_dtype(self, data: NativeT, target_dtype: DType) -> NativeT: def dtype_name(self, dtype: DType) -> str: """The dtype's name as NumPy spells it (``"uint16"``, ``"float32"``).""" - return str(_numpy().dtype(dtype).name) + import numpy as np + + return str(np.dtype(dtype).name) def astype(self, data: NativeT, dtype: DType) -> NativeT: with self._array_scope(data): @@ -713,7 +816,7 @@ class HostDevices(ArrayLibrary[NativeT]): def host_devices(self) -> tuple[HostDevice, ...]: return (HOST,) - def _scope(self, device: Device) -> AbstractContextManager[None]: + def _scope(self, device: Device) -> AbstractContextManager[object]: if isinstance(device, HostDevice): return self._host_scope() return super()._scope(device) @@ -732,7 +835,7 @@ def _place(self, data: np.ndarray, device: Device) -> NativeT: return self._wrap_host_array(data) return super()._place(data, device) - def _host_scope(self) -> AbstractContextManager[None]: + def _host_scope(self) -> AbstractContextManager[object]: return nullcontext() @abstractmethod @@ -752,7 +855,7 @@ def accelerator_devices(self) -> tuple[AcceleratorDevice, ...]: return () return tuple(AcceleratorDevice(index) for index in range(self._accelerator_count())) - def _scope(self, device: Device) -> AbstractContextManager[None]: + def _scope(self, device: Device) -> AbstractContextManager[object]: if isinstance(device, AcceleratorDevice): return self._accelerator_scope(device) return super()._scope(device) @@ -778,7 +881,7 @@ def _accelerator_count(self) -> int: """How many accelerators the imported library sees.""" @abstractmethod - def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[object]: """A block in which new arrays are created on ``device``.""" def _activate_accelerator(self, device: AcceleratorDevice) -> None: @@ -800,10 +903,15 @@ def _argument_device(self, arguments: Sequence[object]) -> Device | None: if self.loaded_module() is None: return None for value in arguments: - if self.accepts(value): - return self.device_of(value) + device = self._device_if_accepted(value) + if device is not None: + return device return None + @abstractmethod + def _device_if_accepted(self, value: object) -> Device | None: + """The device holding ``value``; ``None`` when it is not an array this library takes.""" + def _selected_accelerator(self) -> AcceleratorDevice | None: """The accelerator the library has selected for this thread, where it says.""" @@ -811,7 +919,7 @@ def _selected_accelerator(self) -> AcceleratorDevice | None: def _accelerator_call_scope( self, device: AcceleratorDevice | None - ) -> AbstractContextManager[None]: + ) -> AbstractContextManager[object]: """What a call running on ``device`` is wrapped in.""" del device @@ -855,13 +963,12 @@ class CudaLibrary(AcceleratorDevices[NativeT]): _thread_streams: ClassVar[threading.local] = threading.local() + @abstractmethod def create_stream(self) -> object: """A new CUDA stream on the current device.""" - return self.module.cuda.Stream() - @abstractmethod - def stream_scope(self, stream: object) -> AbstractContextManager[None]: + def stream_scope(self, stream: object) -> AbstractContextManager[object]: """A block whose work is queued on ``stream``.""" @abstractmethod @@ -898,7 +1005,7 @@ def _selected_accelerator(self) -> AcceleratorDevice | None: def _accelerator_call_scope( self, device: AcceleratorDevice | None - ) -> AbstractContextManager[None]: + ) -> AbstractContextManager[object]: if device is None: return nullcontext() return self.stream_scope(self.thread_stream(device)) @@ -921,11 +1028,17 @@ def imports_dlpack_from(self, device: Device) -> bool: # A library without main-memory arrays cannot wrap a main-memory buffer. return isinstance(self, HostDevices) or not isinstance(device, HostDevice) - def from_dlpack(self, data: DLPackPayload | object) -> NativeT: + def from_dlpack(self, data: DLPackPayload | object) -> NativeT | NotImplementedType: """Wrap an exported buffer; ``NotImplemented`` when this capsule cannot be read.""" payload = data if isinstance(data, DLPackPayload) else DLPackPayload(data, data) - return self._import_dlpack(payload) + return self._import_dlpack(payload) if self._reads_dlpack(payload) else NotImplemented + + def _reads_dlpack(self, payload: DLPackPayload) -> bool: + """Whether the library's importer can read this payload.""" + + del payload + return True @abstractmethod def _import_dlpack(self, payload: DLPackPayload) -> NativeT: @@ -1025,84 +1138,95 @@ def floor(self, value: float) -> NativeT: """Floor a host scalar in the library's default precision.""" -class NumpyArrayApi(ArrayMath[NativeT]): +class NumpyArrayApi(ArrayMath[NumpyNativeT]): """Libraries whose array module and arrays follow the NumPy API (NumPy, CuPy, JAX).""" + @property + @abstractmethod + def array_module(self) -> NumpyFunctions[NumpyNativeT]: + """The library's module of NumPy-named functions.""" + def accepts(self, value: object) -> bool: return isinstance(value, self.array_module.ndarray) - def max(self, data: NativeT, axis: Axis) -> NativeT: + def max(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: return self.array_module.max(data, axis=axis) - def min(self, data: NativeT, axis: Axis) -> NativeT: + def min(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: return self.array_module.min(data, axis=axis) - def mean(self, data: NativeT, axis: Axis) -> NativeT: + def mean(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: return self.array_module.mean(data, axis=axis) - def sum(self, data: NativeT, axis: Axis) -> NativeT: + def sum(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: return self.array_module.sum(data, axis=axis) - def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NativeT: + def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> NumpyNativeT: return self.array_module.linspace(start, stop, num, endpoint=endpoint) - def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NativeT: - return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module)) + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> NumpyNativeT: + return self.array_module.ones(shape, dtype=_mapped_dtype(dtype, self.array_module, object)) - def outer(self, left: NativeT, right: NativeT) -> NativeT: + def outer(self, left: NumpyNativeT, right: NumpyNativeT) -> NumpyNativeT: return self.array_module.outer(left, right) - def floor(self, value: float) -> NativeT: + def floor(self, value: float) -> NumpyNativeT: return self.array_module.floor(value) - def _stack(self, values: Sequence[NativeT]) -> NativeT: + def _stack(self, values: Sequence[NumpyNativeT]) -> NumpyNativeT: return self.array_module.stack(tuple(values), axis=0) - def _cast(self, data: NativeT, dtype: DType) -> NativeT: - return getattr(data, "astype")(dtype, copy=False) + def _cast(self, data: NumpyNativeT, dtype: DType) -> NumpyNativeT: + return data.astype(dtype, copy=False) - def _logical_and(self, left: NativeT, right: NativeT) -> NativeT: + def _logical_and(self, left: NumpyNativeT, right: NumpyNativeT) -> NumpyNativeT: return self.array_module.logical_and(left, right) - def _reshape(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: - return getattr(data, "reshape")(shape) + def _reshape(self, data: NumpyNativeT, shape: tuple[int, ...]) -> NumpyNativeT: + return data.reshape(shape) - def _broadcast_to(self, data: NativeT, shape: tuple[int, ...]) -> NativeT: + def _broadcast_to(self, data: NumpyNativeT, shape: tuple[int, ...]) -> NumpyNativeT: return self.array_module.broadcast_to(data, shape) - def _ones_like(self, reference: NativeT, shape: tuple[int, ...], dtype: DType) -> NativeT: + def _ones_like( + self, reference: NumpyNativeT, shape: tuple[int, ...], dtype: DType + ) -> NumpyNativeT: del reference return self.array_module.ones(shape, dtype=dtype) -class MutableArrays(NumpyArrayApi[NativeT], InPlaceAssignment[NativeT]): +class MutableArrays(NumpyArrayApi[MutableNativeT], InPlaceAssignment[MutableNativeT]): """Libraries with the NumPy array API whose arrays can be changed in place.""" + @property + @abstractmethod + def array_module(self) -> MutableNumpyFunctions[MutableNativeT]: + """The library's module of NumPy-named functions.""" + def _normalize_planes( - self, data: NativeT, dtype: DType, scales: Sequence[float | None] - ) -> NativeT: - np = _numpy() + self, data: MutableNativeT, dtype: DType, scales: Sequence[float | None] + ) -> MutableNativeT: + import numpy as np + # Integer division promotes the output dtype, so only floats divide in place. if not np.issubdtype(np.dtype(dtype), np.inexact) or len(scales) == 0: return super()._normalize_planes(data, dtype, scales) - normalized = self.module.array(data, dtype=dtype, copy=True) + normalized = self.array_module.array(data, dtype=dtype, copy=True) for index, scale in enumerate(scales): if scale is not None: normalized[index] /= float(scale) return normalized - def _scale_dtype(self, result: NativeT, target_dtype: DType) -> NativeT: - module = self.module - values = getattr(result, "astype") + def _scale_dtype(self, result: MutableNativeT, target_dtype: DType) -> MutableNativeT: + module = self.array_module if not ( - module.issubdtype(getattr(result, "dtype"), module.floating) - and self._rescales_into(target_dtype) + module.issubdtype(result.dtype, module.floating) and self._rescales_into(target_dtype) ): - return values(target_dtype) + return result.astype(target_dtype) result_min = module.min(result) result_max = module.max(result) if result_max <= result_min: - return values(target_dtype) + 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: @@ -1124,12 +1248,19 @@ class NumpyLibrary(HostDevices["np.ndarray"], MutableArrays["np.ndarray"]): display_name = "NumPy" oom_messages = ("cannot allocate memory", "memory exhausted") + if TYPE_CHECKING: + module = np + + @property + def array_module(self) -> MutableNumpyFunctions[np.ndarray]: + return self.module + def device_of(self, data: np.ndarray) -> Device: del data return HOST def accepts(self, value: object) -> bool: - return isinstance(value, (_numpy().ndarray, ArrayPayload)) + return isinstance(value, (self.module.ndarray, ArrayPayload)) def _to_host(self, data: np.ndarray) -> np.ndarray: return data @@ -1159,13 +1290,23 @@ class CupyLibrary( ("cuda", "runtime", "CUDARuntimeError"), ) + if TYPE_CHECKING: + module = cupy + + @property + def array_module(self) -> MutableNumpyFunctions[cupy.ndarray]: + return self.module + def device_of(self, data: cupy.ndarray) -> Device: return AcceleratorDevice(int(data.device.id)) + def _device_if_accepted(self, value: object) -> Device | None: + return self.device_of(value) if isinstance(value, self.module.ndarray) else None + def _accelerator_count(self) -> int: return int(self.module.cuda.runtime.getDeviceCount()) - def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[object]: return self.module.cuda.Device(device.index) def _activate_accelerator(self, device: AcceleratorDevice) -> None: @@ -1182,23 +1323,28 @@ def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> cupy.ndarray: def current_device(self) -> AcceleratorDevice | None: return AcceleratorDevice(int(self.module.cuda.runtime.getDevice())) - def stream_scope(self, stream: object) -> AbstractContextManager[None]: + def create_stream(self) -> object: + return self.module.cuda.Stream() + + def stream_scope(self, stream: object) -> AbstractContextManager[object]: if not isinstance(stream, AbstractContextManager): raise TypeError(f"{stream!r} is not a CuPy stream") # A CuPy stream is its own context manager. return stream def _free_cached_memory(self) -> None: - module = self.module - module.get_default_memory_pool().free_all_blocks() - module.get_default_pinned_memory_pool().free_all_blocks() - module.cuda.runtime.deviceSynchronize() + self.module.get_default_memory_pool().free_all_blocks() + self.module.get_default_pinned_memory_pool().free_all_blocks() + self.module.cuda.runtime.deviceSynchronize() + + def _reads_dlpack(self, payload: DLPackPayload) -> bool: + # CuPy releases before ``from_dlpack`` took capsules read them with ``fromDlpack``. + return _has_dlpack_protocol(payload.source) or hasattr(self.module, "fromDlpack") def _import_dlpack(self, payload: DLPackPayload) -> cupy.ndarray: if _has_dlpack_protocol(payload.source): return self.module.from_dlpack(payload.source) - capsule_importer = getattr(self.module, "fromDlpack", None) - return NotImplemented if capsule_importer is None else capsule_importer(payload.capsule) + return self.module.fromDlpack(payload.capsule) def _to_host(self, data: cupy.ndarray) -> np.ndarray: return data.get() @@ -1220,14 +1366,19 @@ class TorchLibrary( oom_messages = ("out of memory", "cuda_error_out_of_memory") oom_exception_paths = (("cuda", "OutOfMemoryError"),) + if TYPE_CHECKING: + module = torch + def device_of(self, data: torch.Tensor) -> Device: return AcceleratorDevice(int(data.device.index)) if data.is_cuda else HOST + def _device_if_accepted(self, value: object) -> Device | None: + return self.device_of(value) if isinstance(value, self.module.Tensor) else None + def _accelerator_count(self) -> int: - cuda = self.module.cuda - return int(cuda.device_count()) if cuda.is_available() else 0 + return int(self.module.cuda.device_count()) if self.module.cuda.is_available() else 0 - def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[object]: return self.module.cuda.device(device.index) def _activate_accelerator(self, device: AcceleratorDevice) -> None: @@ -1240,8 +1391,10 @@ def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> torch.Tensor: return self._wrap_host_array(data).to(f"cuda:{device.index}") def _wrap_host_array(self, data: np.ndarray) -> torch.Tensor: + import numpy as np + host_data = ( - _numpy().ascontiguousarray(data) + np.ascontiguousarray(data) if any(stride < 0 for stride in getattr(data, "strides", ())) else data ) @@ -1251,10 +1404,18 @@ def _move_to_host(self, data: torch.Tensor) -> torch.Tensor: return data.cpu() def current_device(self) -> AcceleratorDevice | None: - cuda = self.module.cuda - return AcceleratorDevice(int(cuda.current_device())) if cuda.is_available() else None + return ( + AcceleratorDevice(int(self.module.cuda.current_device())) + if self.module.cuda.is_available() + else None + ) + + def create_stream(self) -> object: + return self.module.cuda.Stream() - def stream_scope(self, stream: object) -> AbstractContextManager[None]: + def stream_scope(self, stream: object) -> AbstractContextManager[object]: + if stream is not None and not isinstance(stream, self.module.cuda.Stream): + raise TypeError(f"{stream!r} is not a PyTorch stream") return self.module.cuda.stream(stream) def _free_cached_memory(self) -> None: @@ -1289,13 +1450,12 @@ def sum(self, data: torch.Tensor, axis: Axis) -> torch.Tensor: return self.module.sum(data, dim=axis) def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> torch.Tensor: - 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] + return self.module.linspace(start, stop, num, dtype=self.module.float32) + return self.module.linspace(start, stop, num + 1, dtype=self.module.float32)[:-1] def ones(self, shape: int | tuple[int, ...], dtype: DType) -> torch.Tensor: - return self.module.ones(shape, dtype=_mapped_dtype(dtype, self.module)) + return self.module.ones(shape, dtype=_mapped_dtype(dtype, self.module, self.module.dtype)) def outer(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: return self.module.outer(left, right) @@ -1307,9 +1467,10 @@ def _to_host(self, data: torch.Tensor) -> np.ndarray: return data.cpu().numpy() def _scale_dtype(self, result: torch.Tensor, target_dtype: DType) -> torch.Tensor: - module, np = self.module, _numpy() - mapped = _mapped_dtype(target_dtype, module) - floats = (module.float16, module.float32, module.float64) + import numpy as np + + mapped = _mapped_dtype(target_dtype, self.module, self.module.dtype) + floats = (self.module.float16, self.module.float32, self.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() @@ -1319,15 +1480,18 @@ def _scale_dtype(self, result: torch.Tensor, target_dtype: DType) -> torch.Tenso 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]) + scaled = self.module.clamp(scaled, min=bounds[0], max=bounds[1]) return scaled.to(mapped) def _stack(self, values: Sequence[torch.Tensor]) -> torch.Tensor: return self.module.stack(tuple(values), dim=0) def _cast(self, data: torch.Tensor, dtype: DType) -> torch.Tensor: - module = self.module - native = dtype if isinstance(dtype, module.dtype) else _mapped_dtype(dtype, module) + native = ( + dtype + if isinstance(dtype, self.module.dtype) + else _mapped_dtype(dtype, self.module, self.module.dtype) + ) return data.to(dtype=native) def _logical_and(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: @@ -1342,8 +1506,11 @@ def _broadcast_to(self, data: torch.Tensor, shape: tuple[int, ...]) -> torch.Ten def _ones_like( self, reference: torch.Tensor, shape: tuple[int, ...], dtype: DType ) -> torch.Tensor: - module = self.module - return module.ones(shape, dtype=_mapped_dtype(dtype, module), device=reference.device) + return self.module.ones( + shape, + dtype=_mapped_dtype(dtype, self.module, self.module.dtype), + device=reference.device, + ) class TensorflowLibrary( @@ -1358,19 +1525,25 @@ class TensorflowLibrary( oom_messages = ("out of memory", "resource_exhausted") oom_exception_paths = (("errors", "ResourceExhaustedError"),) + if TYPE_CHECKING: + module = tensorflow + def device_of(self, data: tensorflow.Tensor) -> Device: device = data.device.lower() if "gpu" not in device: return HOST return AcceleratorDevice(int(device.rsplit(":", maxsplit=1)[-1])) + def _device_if_accepted(self, value: object) -> Device | None: + return self.device_of(value) if isinstance(value, self.module.Tensor) else None + def _accelerator_count(self) -> int: return len(self.module.config.list_logical_devices("GPU")) - def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[object]: return self.module.device(f"/device:GPU:{device.index}") - def _host_scope(self) -> AbstractContextManager[None]: + def _host_scope(self) -> AbstractContextManager[object]: return self.module.device("/device:CPU:0") def _move_to_accelerator( @@ -1394,16 +1567,15 @@ def exports_dlpack_from(self, device: Device) -> bool: return isinstance(device, AcceleratorDevice) def exports_dlpack_data(self, data: tensorflow.Tensor) -> bool: - module = self.module try: - major, minor = map(int, module.__version__.split(".")[:2]) + major, minor = map(int, self.module.__version__.split(".")[:2]) except (AttributeError, TypeError, ValueError): return False if (major, minor) < (2, 12): return False if "gpu" not in str(getattr(data, "device", "")).lower(): return False - dlpack = getattr(getattr(module, "experimental", None), "dlpack", None) + dlpack = getattr(getattr(self.module, "experimental", None), "dlpack", None) return callable(getattr(dlpack, "to_dlpack", None)) def export_dlpack(self, data: tensorflow.Tensor) -> DLPackPayload | None: @@ -1421,7 +1593,9 @@ def accepts(self, value: object) -> bool: return isinstance(value, self.module.Tensor) def dtype_name(self, dtype: DType) -> str: - return str(_numpy().dtype(getattr(dtype, "as_numpy_dtype", dtype)).name) + import numpy as np + + return str(np.dtype(getattr(dtype, "as_numpy_dtype", dtype)).name) def max(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: return self.module.reduce_max(data, axis=axis) @@ -1430,8 +1604,7 @@ def min(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: return self.module.reduce_min(data, axis=axis) def mean(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: - module = self.module - return module.reduce_mean(module.cast(data, module.float32), axis=axis) + return self.module.reduce_mean(self.module.cast(data, self.module.float32), axis=axis) def sum(self, data: tensorflow.Tensor, axis: Axis) -> tensorflow.Tensor: return self.module.reduce_sum(data, axis=axis) @@ -1442,14 +1615,17 @@ def linspace(self, start: float, stop: float, num: int, *, endpoint: bool) -> te return self.module.linspace(float(start), float(stop), num + 1)[:-1] def ones(self, shape: int | tuple[int, ...], dtype: DType) -> tensorflow.Tensor: - return self.module.ones(shape, dtype=_mapped_dtype(dtype, self.module)) + return self.module.ones( + shape, dtype=_mapped_dtype(dtype, self.module, self.module.dtypes.DType) + ) def assign( self, data: tensorflow.Tensor, index: slice, values: tensorflow.Tensor ) -> tensorflow.Tensor: - module = self.module - rows = module.range(*index.indices(int(data.shape[0]))) - return module.tensor_scatter_nd_update(data, module.reshape(rows, [-1, 1]), values) + 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: tensorflow.Tensor, right: tensorflow.Tensor) -> tensorflow.Tensor: return self.module.tensordot(left, right, axes=0) @@ -1461,25 +1637,29 @@ def _to_host(self, data: tensorflow.Tensor) -> np.ndarray: return data.numpy() def _scale_dtype(self, result: tensorflow.Tensor, target_dtype: DType) -> tensorflow.Tensor: - module, np = self.module, _numpy() - mapped = _mapped_dtype(target_dtype, module) - floats = (module.float16, module.float32, module.float64) + import numpy as np + + mapped = _mapped_dtype(target_dtype, self.module, self.module.dtypes.DType) + floats = (self.module.float16, self.module.float32, self.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) + return self.module.cast(result, mapped) + result_min = self.module.reduce_min(result) + result_max = self.module.reduce_max(result) if result_max <= result_min: - return module.cast(result, mapped) + return self.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) + scaled = self.module.clip_by_value(scaled, *bounds) + return self.module.cast(scaled, mapped) def _cast(self, data: tensorflow.Tensor, dtype: DType) -> tensorflow.Tensor: - module = self.module - native = dtype if isinstance(dtype, module.dtypes.DType) else _mapped_dtype(dtype, module) - return module.cast(data, native) + native = ( + dtype + if isinstance(dtype, self.module.dtypes.DType) + else _mapped_dtype(dtype, self.module, self.module.dtypes.DType) + ) + return self.module.cast(data, native) def _stack(self, values: Sequence[tensorflow.Tensor]) -> tensorflow.Tensor: return self.module.stack(tuple(values), axis=0) @@ -1496,9 +1676,10 @@ def _broadcast_to(self, data: tensorflow.Tensor, shape: tuple[int, ...]) -> tens def _ones_like( self, reference: tensorflow.Tensor, shape: tuple[int, ...], dtype: DType ) -> tensorflow.Tensor: - module = self.module - with module.device(reference.device): - return module.ones(shape, dtype=_mapped_dtype(dtype, module)) + with self.module.device(reference.device): + return self.module.ones( + shape, dtype=_mapped_dtype(dtype, self.module, self.module.dtypes.DType) + ) class JaxLibrary( @@ -1514,11 +1695,14 @@ class JaxLibrary( import_environment = (("XLA_PYTHON_CLIENT_PREALLOCATE", "false"),) oom_messages = ("out of memory", "oom when allocating", "allocation failure") + if TYPE_CHECKING: + module = jax + @property - def array_module(self) -> ModuleType: + def array_module(self) -> NumpyFunctions[jax.Array]: return self.module.numpy - def _platform_devices(self, platform: str) -> tuple[object, ...]: + def _platform_devices(self, platform: str) -> tuple[jax.Device, ...]: return tuple(device for device in self.module.devices() if device.platform == platform) def device_of(self, data: jax.Array) -> Device: @@ -1535,13 +1719,16 @@ def device_of(self, data: jax.Array) -> Device: return AcceleratorDevice(index) return NamedDevice(str(device)) + def _device_if_accepted(self, value: object) -> Device | None: + return self.device_of(value) if isinstance(value, self.module.Array) else None + def _accelerator_count(self) -> int: return len(self._platform_devices("gpu")) - def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[None]: + def _accelerator_scope(self, device: AcceleratorDevice) -> AbstractContextManager[object]: return self.module.default_device(self._platform_devices("gpu")[device.index]) - def _host_scope(self) -> AbstractContextManager[None]: + def _host_scope(self) -> AbstractContextManager[object]: return self.module.default_device(self.module.devices("cpu")[0]) def _move_to_accelerator(self, data: jax.Array, device: AcceleratorDevice) -> jax.Array: @@ -1556,9 +1743,10 @@ def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> jax.Array: def _wrap_host_array(self, data: np.ndarray) -> jax.Array: return self.module.device_put(data, self.module.devices("cpu")[0]) + def _reads_dlpack(self, payload: DLPackPayload) -> bool: + return _has_dlpack_protocol(payload.source) + def _import_dlpack(self, payload: DLPackPayload) -> jax.Array: - if not _has_dlpack_protocol(payload.source): - return NotImplemented return self.module.dlpack.from_dlpack(self._aligned(payload.source)) @staticmethod @@ -1582,42 +1770,43 @@ def _aligned(data: object) -> object: return data def mean(self, data: jax.Array, axis: Axis) -> jax.Array: - jnp = self.module.numpy - return jnp.mean(data.astype(jnp.float32), axis=axis) + return self.module.numpy.mean(data.astype(self.module.numpy.float32), axis=axis) def assign(self, data: jax.Array, index: slice, values: jax.Array) -> jax.Array: return data.at[index].set(values) def _to_host(self, data: jax.Array) -> np.ndarray: - return _numpy().asarray(data) + import numpy as np + + return np.asarray(data) def _scale_dtype(self, result: jax.Array, target_dtype: DType) -> jax.Array: - module, np = self.module, _numpy() + import numpy as np + if np.dtype(target_dtype) == np.dtype(np.float64): - x64_enabled = getattr(module.config, "x64_enabled", None) + x64_enabled = getattr(self.module.config, "x64_enabled", None) if x64_enabled is None: - x64_enabled = module.config.read("jax_enable_x64") + x64_enabled = self.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) + mapped = _mapped_dtype(target_dtype, self.module.numpy, object) + floats = (self.module.numpy.float16, self.module.numpy.float32, self.module.numpy.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) + result_min = self.module.numpy.min(result) + result_max = self.module.numpy.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) + scaled = self.module.numpy.clip(scaled, *bounds) return scaled.astype(mapped) def _cast(self, data: jax.Array, dtype: DType) -> jax.Array: - return data.astype(_mapped_dtype(dtype, self.module.numpy)) + return data.astype(_mapped_dtype(dtype, self.module.numpy, object)) def _ones_like(self, reference: jax.Array, shape: tuple[int, ...], dtype: DType) -> jax.Array: return self.module.numpy.ones(shape, dtype=dtype, device=reference.device) @@ -1629,6 +1818,9 @@ class PyclesperantoLibrary( name = "pyclesperanto" oom_messages = ("cl_mem_object_allocation_failure", "cl_out_of_resources", "out of memory") + if TYPE_CHECKING: + module = pyclesperanto + _native_dtype_names: ClassVar[frozenset[str]] = frozenset( {"float32", "int8", "int16", "int32", "uint8", "uint16", "uint32"} ) @@ -1659,12 +1851,16 @@ def _device_named(self, name: str) -> Device: # An OpenCL device that is not a listed GPU, such as a CPU OpenCL device. return NamedDevice(name) - def device_of(self, data: pyclesperanto.Array) -> Device: - declared_device = getattr(data, "device", None) + def device_of(self, data: object) -> Device: + # A host array has no device of its own: a kernel pushes it to the selected one. + declared_device: object = getattr(data, "device", None) if declared_device is None: declared_device = self.module.get_device() return self._device_named(self._device_name(declared_device)) + def _device_if_accepted(self, value: object) -> Device | None: + return self.device_of(value) if self.accepts(value) else None + def _accelerator_count(self) -> int: return len(self._gpu_devices()) @@ -1686,9 +1882,8 @@ def _move_to_accelerator( self, data: pyclesperanto.Array, device: AcceleratorDevice ) -> pyclesperanto.Array: del device - module = self.module - result = module.create_like(data) - module.copy(data, result) + result = self.module.create_like(data) + self.module.copy(data, result) return result def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> pyclesperanto.Array: @@ -1731,56 +1926,56 @@ def _to_host(self, data: pyclesperanto.Array) -> np.ndarray: return self.module.pull(data) def _scale_dtype(self, result: pyclesperanto.Array, target_dtype: DType) -> pyclesperanto.Array: - module, np = self.module, _numpy() + import numpy as np + 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)) + return self.module.push(self.module.pull(result).astype(target_dtype)) + result_min = float(self.module.minimum_of_all_pixels(result)) + result_max = float(self.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( + return self.module.push(self.module.pull(result).astype(target_dtype)) + normalized = self.module.subtract_image_from_scalar(result, scalar=result_min) + normalized = self.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) + scaled = self.module.multiply_image_and_scalar(normalized, scalar=scale) + scaled = self.module.subtract_image_from_scalar(scaled, scalar=offset) elif range_info is not None: - scaled = module.multiply_image_and_scalar(normalized, scalar=range_info) + scaled = self.module.multiply_image_and_scalar(normalized, scalar=range_info) else: scaled = normalized - host_values = module.pull(scaled) + host_values = self.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)) + return self.module.push(host_values.astype(target_dtype)) def _stack(self, values: Sequence[pyclesperanto.Array]) -> pyclesperanto.Array: - module = self.module 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 = self.module.create((1, *source.shape), dtype=source.dtype) + return self.module.copy_slice(source, result, 0) result = values[0] for value in values[1:]: - result = module.concatenate_along_z(result, value) + result = self.module.concatenate_along_z(result, value) return result def _cast(self, data: pyclesperanto.Array, dtype: DType) -> pyclesperanto.Array: - module = self.module - return module.push(module.pull(data).astype(dtype, copy=False)) + return self.module.push(self.module.pull(data).astype(dtype, copy=False)) def _logical_and( self, left: pyclesperanto.Array, right: pyclesperanto.Array ) -> pyclesperanto.Array: - module = self.module - return module.push(_numpy().logical_and(module.pull(left), module.pull(right))) + import numpy as np + + return self.module.push(np.logical_and(self.module.pull(left), self.module.pull(right))) def _reshape(self, data: pyclesperanto.Array, shape: tuple[int, ...]) -> pyclesperanto.Array: if tuple(data.shape) == shape: @@ -1799,13 +1994,14 @@ def _broadcast_to( def _ones_like( self, reference: pyclesperanto.Array, shape: tuple[int, ...], dtype: DType ) -> pyclesperanto.Array: - module = self.module - if _numpy().dtype(dtype).name not in self._native_dtype_names: + import numpy as np + + if np.dtype(dtype).name not in self._native_dtype_names: 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" ) - result = module.create(shape, dtype=dtype, device=reference.device) - module.set(result, scalar=1, device=reference.device) + result = self.module.create(shape, dtype=dtype, device=reference.device) + self.module.set(result, scalar=1, device=reference.device) return result diff --git a/src/arraybridge/types.py b/src/arraybridge/types.py index a317bc6..88890e1 100644 --- a/src/arraybridge/types.py +++ b/src/arraybridge/types.py @@ -12,70 +12,98 @@ from __future__ import annotations -from collections.abc import Mapping, MutableMapping +import sys +from collections.abc import Iterator, Mapping, MutableMapping from enum import Enum, EnumMeta from types import MappingProxyType +from typing import TypeGuard, TypeVar from arraybridge.libraries import Array, ArrayLibrary +MemberT = TypeVar("MemberT") + + +def _add_missing_members(enum_class: type) -> None: + """Enter a member for every registered library name the enum does not have yet.""" + + if not issubclass(enum_class, Enum): + raise TypeError(f"{enum_class!r} is not an enum class") + for name in ArrayLibrary.__registry__: + if name not in enum_class._value2member_map_: + member = object.__new__(enum_class) + member._name_ = name.upper() + member._value_ = name + enum_class._member_map_[member._name_] = member + enum_class._member_names_.append(member._name_) + enum_class._value2member_map_[name] = member + + +def _holds_members_of( + members: MappingProxyType[str, Enum], enum_class: type[MemberT] +) -> TypeGuard[MappingProxyType[str, MemberT]]: + return all(isinstance(member, enum_class) for member in members.values()) + class _RegisteredLibraryNames(EnumMeta): """Keeps the members equal to the registered library names. The standard library offers no way to add a member to an existing enum, so - a missing one is entered into the enum's own member tables here. The - overrides below keep ``EnumMeta``'s signatures and add only that step. + a missing one is entered into the enum's own member tables. The overrides + below keep ``EnumMeta``'s signatures and add only that step. A lookup by + value needs no override: ``Enum`` asks ``_missing_`` for a value it does + not know, and ``MemoryType._missing_`` adds the missing members. """ - def _add_missing_members(cls): - for name in ArrayLibrary.__registry__: - if name not in cls._value2member_map_: - member = object.__new__(cls) - member._name_ = name.upper() - member._value_ = name - cls._member_map_[member._name_] = member - cls._member_names_.append(member._name_) - cls._value2member_map_[name] = member - - def __call__(cls, *args, **kwargs): - cls._add_missing_members() - return super().__call__(*args, **kwargs) - - def __getitem__(cls, name): - cls._add_missing_members() - return super().__getitem__(name) - - def __reversed__(cls): - cls._add_missing_members() - return super().__reversed__() - - def __dir__(cls): - cls._add_missing_members() + def __getitem__(cls: type[MemberT], name: str) -> MemberT: + _add_missing_members(cls) + return EnumMeta.__getitem__(cls, name) + + def __reversed__(cls: type[MemberT]) -> Iterator[MemberT]: + _add_missing_members(cls) + return EnumMeta.__reversed__(cls) + + def __dir__(cls) -> list[str]: + _add_missing_members(cls) return super().__dir__() - def __iter__(cls): - cls._add_missing_members() - return super().__iter__() + def __iter__(cls: type[MemberT]) -> Iterator[MemberT]: + _add_missing_members(cls) + return EnumMeta.__iter__(cls) - def __len__(cls): - cls._add_missing_members() + def __len__(cls) -> int: + _add_missing_members(cls) return super().__len__() - def __contains__(cls, value): - cls._add_missing_members() - return super().__contains__(value) + # ``EnumMeta.__contains__`` takes any value from Python 3.12 and only a member before. + if sys.version_info >= (3, 12): - @property - def __members__(cls): - cls._add_missing_members() - return MappingProxyType(cls._member_map_) + def __contains__(cls, value: object) -> bool: + _add_missing_members(cls) + return super().__contains__(value) - def __getattr__(cls, name): + else: + + def __contains__(cls, member: Enum) -> bool: + _add_missing_members(cls) + return super().__contains__(member) + + @property + def __members__(cls: type[MemberT]) -> MappingProxyType[str, MemberT]: + _add_missing_members(cls) + if not issubclass(cls, Enum): + raise TypeError(f"{cls!r} is not an enum class") + members = MappingProxyType(cls._member_map_) + if not _holds_members_of(members, cls): + raise TypeError(f"{cls!r} holds a member of another class") + return members + + def __getattr__(cls: type[MemberT], name: str) -> MemberT: if not name.startswith("_"): - cls._add_missing_members() - member = cls._member_map_.get(name) - if member is not None: - return member + _add_missing_members(cls) + try: + return EnumMeta.__getitem__(cls, name) + except KeyError: + pass raise AttributeError(name) @@ -84,7 +112,7 @@ class MemoryType(Enum, metaclass=_RegisteredLibraryNames): @classmethod def _missing_(cls, value: object) -> MemoryType | None: - cls._add_missing_members() + _add_missing_members(cls) member = cls._value2member_map_.get(value) return member if isinstance(member, cls) else None @@ -95,7 +123,7 @@ def library(self) -> ArrayLibrary[Array]: return ArrayLibrary.named(self.value) -MemoryType._add_missing_members() +_add_missing_members(MemoryType) class MemoryContractAttribute(str, Enum): @@ -108,9 +136,14 @@ class MemoryContractAttribute(str, Enum): def read(self, namespace: object, default: str | None = None) -> str | None: """Read this attribute from an object or a mapping.""" - if isinstance(namespace, Mapping): - return namespace.get(self.value, default) - return getattr(namespace, self.value, default) + value: object = ( + namespace.get(self.value, default) + if isinstance(namespace, Mapping) + else getattr(namespace, self.value, default) + ) + if value is None or isinstance(value, str): + return value + raise TypeError(f"{self.value} of {namespace!r} is {value!r}, not a library name") def write(self, namespace: object, value: str) -> None: """Write this attribute to an object or a mutable mapping.""" diff --git a/stubs/cupy/__init__.pyi b/stubs/cupy/__init__.pyi new file mode 100644 index 0000000..b6d6f8d --- /dev/null +++ b/stubs/cupy/__init__.pyi @@ -0,0 +1,52 @@ +"""What arraybridge calls of CuPy.""" + +from collections.abc import Sequence + +import numpy +from typing_extensions import Self + +from . import cuda as cuda + +_Axis = int | tuple[int, ...] + +class floating: ... + +class ndarray: + shape: tuple[int, ...] + dtype: numpy.dtype + device: cuda.Device + def astype(self, dtype: object, *, copy: bool = ...) -> Self: ... + def reshape(self, shape: tuple[int, ...]) -> Self: ... + def copy(self) -> Self: ... + def get(self) -> numpy.ndarray: ... + def __getitem__(self, index: int | slice) -> Self: ... + def __setitem__(self, index: int | slice, value: Self) -> None: ... + def __sub__(self, other: Self | float) -> Self: ... + def __mul__(self, other: float) -> Self: ... + def __truediv__(self, other: Self | float) -> Self: ... + def __itruediv__(self, other: Self | float) -> Self: ... + def __le__(self, other: Self) -> Self: ... + +class _MemoryPool: + def free_all_blocks(self) -> None: ... + +def get_default_memory_pool() -> _MemoryPool: ... +def get_default_pinned_memory_pool() -> _MemoryPool: ... +def from_dlpack(array: object) -> ndarray: ... + +# Present only in CuPy releases that still read bare capsules. +def fromDlpack(dltensor: object) -> ndarray: ... +def issubdtype(arg1: object, arg2: object) -> bool: ... +def array(a: numpy.ndarray | ndarray, *, dtype: object = ..., copy: bool = ...) -> ndarray: ... +def clip(a: ndarray, a_min: float, a_max: float) -> ndarray: ... +def max(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def min(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def mean(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def sum(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def linspace(start: float, stop: float, num: int, *, endpoint: bool = ...) -> ndarray: ... +def ones(shape: int | tuple[int, ...], dtype: object = ...) -> ndarray: ... +def outer(a: ndarray, b: ndarray) -> ndarray: ... +def floor(x: float) -> ndarray: ... +def stack(arrays: Sequence[ndarray], axis: int = ...) -> ndarray: ... +def logical_and(x1: ndarray, x2: ndarray) -> ndarray: ... +def broadcast_to(array: ndarray, shape: tuple[int, ...]) -> ndarray: ... diff --git a/stubs/cupy/cuda/__init__.pyi b/stubs/cupy/cuda/__init__.pyi new file mode 100644 index 0000000..1b5bf74 --- /dev/null +++ b/stubs/cupy/cuda/__init__.pyi @@ -0,0 +1,29 @@ +"""What arraybridge calls of ``cupy.cuda``.""" + +from types import TracebackType + +from typing_extensions import Self + +from . import runtime as runtime + +class Device: + id: int + def __init__(self, device: int = ...) -> None: ... + def use(self) -> None: ... + def __enter__(self) -> Self: ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: ... + +class Stream: + def __init__(self) -> None: ... + def __enter__(self) -> Self: ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: ... diff --git a/stubs/cupy/cuda/runtime.pyi b/stubs/cupy/cuda/runtime.pyi new file mode 100644 index 0000000..15aaaf2 --- /dev/null +++ b/stubs/cupy/cuda/runtime.pyi @@ -0,0 +1,5 @@ +"""What arraybridge calls of ``cupy.cuda.runtime``.""" + +def getDeviceCount() -> int: ... +def getDevice() -> int: ... +def deviceSynchronize() -> None: ... diff --git a/stubs/jax/__init__.pyi b/stubs/jax/__init__.pyi new file mode 100644 index 0000000..48adf99 --- /dev/null +++ b/stubs/jax/__init__.pyi @@ -0,0 +1,41 @@ +"""What arraybridge calls of JAX.""" + +from collections.abc import Callable +from contextlib import AbstractContextManager + +import numpy as _host_numpy +from typing_extensions import Self + +from . import dlpack as dlpack +from . import numpy as numpy + +class Device: + platform: str + +class _IndexUpdate: + def set(self, values: Array) -> Array: ... + +class _IndexUpdates: + def __getitem__(self, index: slice) -> _IndexUpdate: ... + +class Array: + shape: tuple[int, ...] + dtype: _host_numpy.dtype + at: _IndexUpdates + # A property in current JAX, a method in older releases. + device: Device | Callable[[], Device] + def astype(self, dtype: object, *, copy: bool = ...) -> Self: ... + def reshape(self, shape: tuple[int, ...]) -> Self: ... + def __sub__(self, other: Self | float) -> Self: ... + def __mul__(self, other: float) -> Self: ... + def __truediv__(self, other: Self | float) -> Self: ... + def __le__(self, other: Self) -> Self: ... + +class _Config: + def read(self, name: str) -> bool: ... + +config: _Config + +def devices(backend: str | None = ...) -> list[Device]: ... +def default_device(new_val: Device) -> AbstractContextManager[None]: ... +def device_put(x: _host_numpy.ndarray | Array, device: Device) -> Array: ... diff --git a/stubs/jax/dlpack.pyi b/stubs/jax/dlpack.pyi new file mode 100644 index 0000000..a34f34d --- /dev/null +++ b/stubs/jax/dlpack.pyi @@ -0,0 +1,5 @@ +"""What arraybridge calls of ``jax.dlpack``.""" + +from jax import Array + +def from_dlpack(external_array: object) -> Array: ... diff --git a/stubs/jax/numpy.pyi b/stubs/jax/numpy.pyi new file mode 100644 index 0000000..fc2fdc2 --- /dev/null +++ b/stubs/jax/numpy.pyi @@ -0,0 +1,33 @@ +"""What arraybridge calls of ``jax.numpy``.""" + +from collections.abc import Callable, Sequence + +from jax import Array, Device + +ndarray = Array + +_Axis = int | tuple[int, ...] + +class _ScalarType: ... + +float16: _ScalarType +float32: _ScalarType +float64: _ScalarType + +def clip(arr: Array, min: float, max: float) -> Array: ... +def max(a: Array, axis: _Axis | None = ...) -> Array: ... +def min(a: Array, axis: _Axis | None = ...) -> Array: ... +def mean(a: Array, axis: _Axis | None = ...) -> Array: ... +def sum(a: Array, axis: _Axis | None = ...) -> Array: ... +def linspace(start: float, stop: float, num: int, *, endpoint: bool = ...) -> Array: ... +def ones( + shape: int | tuple[int, ...], + dtype: object = ..., + *, + device: Device | Callable[[], Device] = ..., +) -> Array: ... +def outer(a: Array, b: Array) -> Array: ... +def floor(x: float) -> Array: ... +def stack(arrays: Sequence[Array], axis: int = ...) -> Array: ... +def logical_and(x1: Array, x2: Array) -> Array: ... +def broadcast_to(array: Array, shape: tuple[int, ...]) -> Array: ... diff --git a/stubs/metaclass_registry/__init__.pyi b/stubs/metaclass_registry/__init__.pyi new file mode 100644 index 0000000..4bb269e --- /dev/null +++ b/stubs/metaclass_registry/__init__.pyi @@ -0,0 +1,12 @@ +"""What arraybridge uses of metaclass-registry.""" + +from abc import ABCMeta +from enum import Enum + +class RegistryKeyAttribute(str, Enum): + STRATEGY_LABEL = "strategy_label" + +class RegistryFamily: + def __init__(self, key_attribute: str | RegistryKeyAttribute) -> None: ... + +class AutoRegisterMeta(ABCMeta): ... diff --git a/stubs/numpy/__init__.pyi b/stubs/numpy/__init__.pyi new file mode 100644 index 0000000..c8db68b --- /dev/null +++ b/stubs/numpy/__init__.pyi @@ -0,0 +1,61 @@ +"""What arraybridge calls of NumPy.""" + +from collections.abc import Sequence + +from typing_extensions import Self + +_Axis = int | tuple[int, ...] + +class dtype: + name: str + itemsize: int + def __init__(self, dtype: object) -> None: ... + +class generic: ... +class number(generic): ... + +class integer(number): + def __int__(self) -> int: ... + +class inexact(number): ... +class floating(inexact): ... +class uint8(integer): ... +class uint16(integer): ... +class int16(integer): ... +class int32(integer): ... +class int64(integer): ... +class float32(floating): ... +class float64(floating): ... + +class ndarray: + shape: tuple[int, ...] + dtype: dtype + nbytes: int + strides: tuple[int, ...] + def astype(self, dtype: object, *, copy: bool = ...) -> Self: ... + def reshape(self, shape: tuple[int, ...]) -> Self: ... + def __getitem__(self, index: int | slice) -> Self: ... + def __setitem__(self, index: int | slice, value: Self) -> None: ... + def __sub__(self, other: Self | float) -> Self: ... + def __mul__(self, other: float) -> Self: ... + def __truediv__(self, other: Self | float) -> Self: ... + def __itruediv__(self, other: Self | float) -> Self: ... + def __le__(self, other: Self) -> Self: ... + +def issubdtype(arg1: object, arg2: object) -> bool: ... +def prod(a: Sequence[int], *, dtype: object = ...) -> int64: ... +def asarray(a: object) -> ndarray: ... +def ascontiguousarray(a: ndarray) -> ndarray: ... +def array(a: ndarray, *, dtype: object = ..., copy: bool = ...) -> ndarray: ... +def clip(a: ndarray, a_min: float, a_max: float) -> ndarray: ... +def max(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def min(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def mean(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def sum(a: ndarray, axis: _Axis | None = ...) -> ndarray: ... +def linspace(start: float, stop: float, num: int, *, endpoint: bool = ...) -> ndarray: ... +def ones(shape: int | tuple[int, ...], dtype: object = ...) -> ndarray: ... +def outer(a: ndarray, b: ndarray) -> ndarray: ... +def floor(x: float) -> ndarray: ... +def stack(arrays: Sequence[ndarray], axis: int = ...) -> ndarray: ... +def logical_and(x1: ndarray, x2: ndarray) -> ndarray: ... +def broadcast_to(array: ndarray, shape: tuple[int, ...]) -> ndarray: ... diff --git a/stubs/pyclesperanto/__init__.pyi b/stubs/pyclesperanto/__init__.pyi new file mode 100644 index 0000000..6951f0f --- /dev/null +++ b/stubs/pyclesperanto/__init__.pyi @@ -0,0 +1,32 @@ +"""What arraybridge calls of pyclesperanto.""" + +import numpy + +class Device: + name: str + +class Array: + shape: tuple[int, ...] + dtype: numpy.dtype + device: Device + def __setitem__(self, index: int, value: Array) -> None: ... + +def list_available_devices(device_type: str = ...) -> list[Device]: ... +def select_device(device_name: str | int = ..., device_type: str = ...) -> Device: ... +def get_device() -> Device: ... +def create(shape: tuple[int, ...], dtype: object = ..., device: Device = ...) -> Array: ... +def create_like(array: Array) -> Array: ... +def copy(input_image: Array, output_image: Array) -> Array: ... +def copy_slice(input_image: Array, output_image: Array, slice_index: int) -> Array: ... +def push(array: numpy.ndarray) -> Array: ... +def pull(array: Array) -> numpy.ndarray: ... +def set(input_image: Array, *, scalar: float, device: Device = ...) -> Array: ... +def concatenate_along_z(input_image0: Array, input_image1: Array) -> Array: ... +def maximum_z_projection(input_image: Array) -> Array: ... +def minimum_z_projection(input_image: Array) -> Array: ... +def mean_z_projection(input_image: Array) -> Array: ... +def sum_z_projection(input_image: Array) -> Array: ... +def minimum_of_all_pixels(input_image: Array) -> float: ... +def maximum_of_all_pixels(input_image: Array) -> float: ... +def subtract_image_from_scalar(input_image: Array, *, scalar: float) -> Array: ... +def multiply_image_and_scalar(input_image: Array, *, scalar: float) -> Array: ... diff --git a/stubs/python_introspect/__init__.pyi b/stubs/python_introspect/__init__.pyi new file mode 100644 index 0000000..d5de834 --- /dev/null +++ b/stubs/python_introspect/__init__.pyi @@ -0,0 +1,15 @@ +"""What arraybridge uses of python-introspect.""" + +import inspect +from abc import ABC, abstractmethod +from typing import ClassVar + +class RuntimeParameterDeclarationABC(ABC): + preserve_for_execution: ClassVar[bool] + is_semantic_control: ClassVar[bool] + @classmethod + @abstractmethod + def require_parameter_name(cls) -> str: ... + @classmethod + @abstractmethod + def parameter(cls) -> inspect.Parameter: ... diff --git a/stubs/tensorflow/__init__.pyi b/stubs/tensorflow/__init__.pyi new file mode 100644 index 0000000..742187a --- /dev/null +++ b/stubs/tensorflow/__init__.pyi @@ -0,0 +1,54 @@ +"""What arraybridge calls of TensorFlow.""" + +from collections.abc import Sequence +from contextlib import AbstractContextManager + +import numpy +from typing_extensions import Self + +from . import config as config +from . import dtypes as dtypes +from . import experimental as experimental +from . import math as math + +_Axis = int | tuple[int, ...] + +__version__: str + +float16: dtypes.DType +float32: dtypes.DType +float64: dtypes.DType + +class TensorShape: + def num_elements(self) -> int: ... + def __getitem__(self, index: int) -> int: ... + +class Tensor: + device: str + shape: TensorShape + dtype: dtypes.DType + def numpy(self) -> numpy.ndarray: ... + def __getitem__(self, index: int | slice) -> Self: ... + def __sub__(self, other: Self | float) -> Self: ... + def __mul__(self, other: float) -> Self: ... + def __truediv__(self, other: Self | float) -> Self: ... + def __le__(self, other: Self) -> Self: ... + +def device(device_name: str) -> AbstractContextManager[None]: ... +def identity(input: Tensor) -> Tensor: ... +def convert_to_tensor(value: numpy.ndarray) -> Tensor: ... +def cast(x: Tensor, dtype: dtypes.DType) -> Tensor: ... +def reduce_max(input_tensor: Tensor, axis: _Axis | None = ...) -> Tensor: ... +def reduce_min(input_tensor: Tensor, axis: _Axis | None = ...) -> Tensor: ... +def reduce_mean(input_tensor: Tensor, axis: _Axis | None = ...) -> Tensor: ... +def reduce_sum(input_tensor: Tensor, axis: _Axis | None = ...) -> Tensor: ... +def linspace(start: float, stop: float, num: int) -> Tensor: ... +def ones(shape: int | tuple[int, ...], dtype: dtypes.DType = ...) -> Tensor: ... +def range(start: int, limit: int, delta: int) -> Tensor: ... +def reshape(tensor: Tensor, shape: Sequence[int]) -> Tensor: ... +def tensor_scatter_nd_update(tensor: Tensor, indices: Tensor, updates: Tensor) -> Tensor: ... +def tensordot(a: Tensor, b: Tensor, axes: int) -> Tensor: ... +def clip_by_value(t: Tensor, clip_value_min: float, clip_value_max: float) -> Tensor: ... +def stack(values: Sequence[Tensor], axis: int = ...) -> Tensor: ... +def logical_and(x: Tensor, y: Tensor) -> Tensor: ... +def broadcast_to(input: Tensor, shape: tuple[int, ...]) -> Tensor: ... diff --git a/stubs/tensorflow/config.pyi b/stubs/tensorflow/config.pyi new file mode 100644 index 0000000..e53f790 --- /dev/null +++ b/stubs/tensorflow/config.pyi @@ -0,0 +1,5 @@ +"""What arraybridge calls of ``tensorflow.config``.""" + +class LogicalDevice: ... + +def list_logical_devices(device_type: str | None = ...) -> list[LogicalDevice]: ... diff --git a/stubs/tensorflow/dtypes.pyi b/stubs/tensorflow/dtypes.pyi new file mode 100644 index 0000000..6701631 --- /dev/null +++ b/stubs/tensorflow/dtypes.pyi @@ -0,0 +1,4 @@ +"""What arraybridge calls of ``tensorflow.dtypes``.""" + +class DType: + size: int diff --git a/stubs/tensorflow/experimental/__init__.pyi b/stubs/tensorflow/experimental/__init__.pyi new file mode 100644 index 0000000..d1417d8 --- /dev/null +++ b/stubs/tensorflow/experimental/__init__.pyi @@ -0,0 +1,3 @@ +"""What arraybridge calls of ``tensorflow.experimental``.""" + +from . import dlpack as dlpack diff --git a/stubs/tensorflow/experimental/dlpack.pyi b/stubs/tensorflow/experimental/dlpack.pyi new file mode 100644 index 0000000..4bbb3ca --- /dev/null +++ b/stubs/tensorflow/experimental/dlpack.pyi @@ -0,0 +1,6 @@ +"""What arraybridge calls of ``tensorflow.experimental.dlpack``.""" + +from tensorflow import Tensor + +def to_dlpack(tf_tensor: Tensor) -> object: ... +def from_dlpack(dlcapsule: object) -> Tensor: ... diff --git a/stubs/tensorflow/math.pyi b/stubs/tensorflow/math.pyi new file mode 100644 index 0000000..c3b866b --- /dev/null +++ b/stubs/tensorflow/math.pyi @@ -0,0 +1,5 @@ +"""What arraybridge calls of ``tensorflow.math``.""" + +from tensorflow import Tensor + +def floor(x: float) -> Tensor: ... diff --git a/stubs/torch/__init__.pyi b/stubs/torch/__init__.pyi new file mode 100644 index 0000000..53d02af --- /dev/null +++ b/stubs/torch/__init__.pyi @@ -0,0 +1,58 @@ +"""What arraybridge calls of PyTorch.""" + +import builtins +from collections.abc import Sequence + +import numpy +from typing_extensions import Self + +from . import cuda as cuda + +_Axis = int | tuple[int, ...] + +class dtype: ... + +_DType = dtype + +uint16: dtype +float16: dtype +float32: dtype +float64: dtype + +class device: + index: int + +class Tensor: + device: device + is_cuda: bool + dtype: _DType + def to(self, device: str | _DType = ..., *, dtype: _DType = ...) -> Self: ... + def cpu(self) -> Self: ... + def float(self) -> Self: ... + def numpy(self) -> numpy.ndarray: ... + def element_size(self) -> int: ... + def nelement(self) -> int: ... + def min(self) -> Self: ... + def max(self) -> Self: ... + def reshape(self, shape: tuple[int, ...]) -> Self: ... + def expand(self, size: tuple[int, ...]) -> Self: ... + def __getitem__(self, index: int | slice) -> Self: ... + def __sub__(self, other: Self | builtins.float) -> Self: ... + def __mul__(self, other: builtins.float) -> Self: ... + def __truediv__(self, other: Self | builtins.float) -> Self: ... + def __le__(self, other: Self) -> Self: ... + +def from_numpy(ndarray: numpy.ndarray) -> Tensor: ... +def from_dlpack(ext_tensor: object) -> Tensor: ... +def tensor(data: float) -> Tensor: ... +def amax(input: Tensor, dim: _Axis) -> Tensor: ... +def amin(input: Tensor, dim: _Axis) -> Tensor: ... +def mean(input: Tensor, dim: _Axis) -> Tensor: ... +def sum(input: Tensor, dim: _Axis) -> Tensor: ... +def linspace(start: float, end: float, steps: int, *, dtype: dtype = ...) -> Tensor: ... +def ones(size: int | tuple[int, ...], *, dtype: dtype = ..., device: device = ...) -> Tensor: ... +def outer(input: Tensor, vec2: Tensor) -> Tensor: ... +def floor(input: Tensor) -> Tensor: ... +def clamp(input: Tensor, *, min: float, max: float) -> Tensor: ... +def stack(tensors: Sequence[Tensor], dim: int = ...) -> Tensor: ... +def logical_and(input: Tensor, other: Tensor) -> Tensor: ... diff --git a/stubs/torch/cuda.pyi b/stubs/torch/cuda.pyi new file mode 100644 index 0000000..ba9b581 --- /dev/null +++ b/stubs/torch/cuda.pyi @@ -0,0 +1,33 @@ +"""What arraybridge calls of ``torch.cuda``.""" + +from types import TracebackType + +class Stream: + def __init__(self) -> None: ... + +class device: + def __init__(self, device: int) -> None: ... + def __enter__(self) -> None: ... + def __exit__( + self, + type: type[BaseException] | None, + value: BaseException | None, + traceback: TracebackType | None, + ) -> None: ... + +class _StreamScope: + def __enter__(self) -> None: ... + def __exit__( + self, + type: type[BaseException] | None, + value: BaseException | None, + traceback: TracebackType | None, + ) -> None: ... + +def stream(stream: Stream | None) -> _StreamScope: ... +def is_available() -> bool: ... +def device_count() -> int: ... +def set_device(device: int) -> None: ... +def current_device() -> int: ... +def empty_cache() -> None: ... +def synchronize() -> None: ... diff --git a/tests/fakes.py b/tests/fakes.py index 22ee7c4..553ea33 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -167,6 +167,14 @@ class FakeTorch: Tensor = FakeTorchTensor def __init__(self, device_count=0): + module = self + + class Stream(FakeStream): + """``torch.cuda.Stream`` is a class: a stream is recognised by being one.""" + + def __init__(self): + super().__init__(module) + self.events = [] self.current = 0 self.cuda = SimpleNamespace( @@ -175,7 +183,7 @@ def __init__(self, device_count=0): current_device=lambda: self.current, set_device=lambda index: self.events.append(("set_device", index)), device=lambda index: DeviceScope(self, index), - Stream=lambda: FakeStream(self), + Stream=Stream, stream=lambda stream: stream, empty_cache=lambda: self.events.append(("empty_cache", self.current)), synchronize=lambda: self.events.append(("synchronize", self.current)), diff --git a/tests/test_checked_values.py b/tests/test_checked_values.py new file mode 100644 index 0000000..0ffa1b2 --- /dev/null +++ b/tests/test_checked_values.py @@ -0,0 +1,145 @@ +"""The places where a value of unknown type is checked before it is given a type.""" + +import sys +from enum import Enum +from types import MappingProxyType, SimpleNamespace + +import numpy as np +import pytest + +from arraybridge.conversion import ( + DeviceMove, + DLPackExchange, + RouteDeclinedError, + _imports_dlpack, + _is_array_of_both, +) +from arraybridge.devices import HOST, AcceleratorDevice +from arraybridge.libraries import ArrayLibrary, _mapped_dtype +from arraybridge.types import ( + MemoryContractAttribute, + MemoryType, + _add_missing_members, + _holds_members_of, +) +from tests.fakes import FakeCupy, FakeCupyArray, FakePyclesperanto, FakeTorch, FakeTorchTensor + +NUMPY = ArrayLibrary.named("numpy") +TORCH = ArrayLibrary.named("torch") +CUPY = ArrayLibrary.named("cupy") +PYCLESPERANTO = ArrayLibrary.named("pyclesperanto") + + +class TestMappedDtype: + def test_returns_the_object_of_the_asked_class(self): + assert _mapped_dtype("float32", np, type) is np.float32 + assert _mapped_dtype(np.uint16, np, object) is np.uint16 + + def test_rejects_an_object_of_another_class(self): + with pytest.raises(TypeError, match="numpy.float32 is not a int"): + _mapped_dtype("float32", np, int) + + def test_rejects_a_name_the_module_does_not_have(self): + with pytest.raises(TypeError, match="does not expose dtype float32"): + _mapped_dtype("float32", SimpleNamespace(), object) + + +class TestMemoryContractAttributeRead: + @pytest.mark.parametrize("namespace", [{}, SimpleNamespace()], ids=["mapping", "object"]) + def test_returns_a_written_name_and_the_default(self, namespace): + assert MemoryContractAttribute.INPUT.read(namespace) is None + assert MemoryContractAttribute.INPUT.read(namespace, "numpy") == "numpy" + MemoryContractAttribute.INPUT.write(namespace, "torch") + assert MemoryContractAttribute.INPUT.read(namespace) == "torch" + + @pytest.mark.parametrize( + "namespace", + [{"input_memory_type": 3}, SimpleNamespace(input_memory_type=3)], + ids=["mapping", "object"], + ) + def test_rejects_a_value_that_is_not_a_name(self, namespace): + with pytest.raises(TypeError, match="not a library name"): + MemoryContractAttribute.INPUT.read(namespace) + + +class TestMemoryTypeMembers: + def test_members_of_the_enum_are_recognised(self): + members = MappingProxyType({"NUMPY": MemoryType.NUMPY}) + assert _holds_members_of(members, MemoryType) + assert MemoryType.__members__["NUMPY"] is MemoryType.NUMPY + + def test_a_member_of_another_enum_is_not(self): + class Other(Enum): + NUMPY = "numpy" + + assert not _holds_members_of(MappingProxyType({"NUMPY": Other.NUMPY}), MemoryType) + + def test_only_an_enum_class_takes_members(self): + with pytest.raises(TypeError, match="not an enum class"): + _add_missing_members(int) + + +class TestDeviceMove: + def test_an_array_of_one_library_is_an_array_of_the_same_library(self): + values = np.arange(3) + assert _is_array_of_both(values, NUMPY, NUMPY) + assert DeviceMove.convert(values, NUMPY, NUMPY, HOST) is values + + def test_two_libraries_are_rejected(self): + values = np.arange(3) + assert not _is_array_of_both(values, NUMPY, TORCH) + with pytest.raises(TypeError, match="different libraries"): + DeviceMove.convert(values, NUMPY, TORCH, HOST) + + +class TestDLPackTarget: + def test_only_a_dlpack_importer_is_one(self): + assert _imports_dlpack(TORCH) + assert not _imports_dlpack(NUMPY) + + def test_a_target_that_does_not_import_declines_the_route(self): + with pytest.raises(RouteDeclinedError, match="does not import DLPack"): + DLPackExchange.convert(np.arange(3), NUMPY, NUMPY, HOST) + + +class TestTorchStreamScope: + def test_takes_a_torch_stream_and_no_stream(self, monkeypatch): + module = FakeTorch(device_count=1) + monkeypatch.setitem(sys.modules, "torch", module) + stream = TORCH.create_stream() + + assert TORCH.stream_scope(stream) is stream + assert TORCH.stream_scope(None) is None + + def test_rejects_what_is_not_a_stream(self, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", FakeTorch(device_count=1)) + + with pytest.raises(TypeError, match="is not a PyTorch stream"): + TORCH.stream_scope(object()) + + +class TestDeviceIfAccepted: + def test_an_array_of_the_library_has_its_device(self, monkeypatch): + cupy = FakeCupy(device_count=2) + monkeypatch.setitem(sys.modules, "cupy", cupy) + monkeypatch.setitem(sys.modules, "torch", FakeTorch(device_count=2)) + + assert CUPY._device_if_accepted(FakeCupyArray([1], 1, cupy)) == AcceleratorDevice(1) + assert TORCH._device_if_accepted(FakeTorchTensor([1], cuda_index=1)) == AcceleratorDevice(1) + assert TORCH._device_if_accepted(FakeTorchTensor([1])) == HOST + + def test_another_object_has_none(self, monkeypatch): + monkeypatch.setitem(sys.modules, "cupy", FakeCupy(device_count=2)) + monkeypatch.setitem(sys.modules, "torch", FakeTorch(device_count=2)) + + for value in (np.arange(3), "text", None): + assert CUPY._device_if_accepted(value) is None + assert TORCH._device_if_accepted(value) is None + + def test_pyclesperanto_takes_a_host_array_on_its_selected_device(self, monkeypatch): + monkeypatch.setitem(sys.modules, "pyclesperanto", FakePyclesperanto()) + + assert PYCLESPERANTO._device_if_accepted(np.arange(3)) == PYCLESPERANTO.device_of( + np.arange(3) + ) + assert PYCLESPERANTO._device_if_accepted("text") is None From 60b96627cf7c26173a73d8caad642822b3996225 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 15:59:16 -0400 Subject: [PATCH 5/5] 0.5.1: changelog, version, GPU commands Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 107 ++++++++++++++++++++++++++++-------- TESTING.md | 35 ++++++++++++ src/arraybridge/__init__.py | 2 +- 3 files changed, 120 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3544e41..9c6639a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,89 @@ # Changelog +## 0.5.1 + +### No library inherits an operation it cannot perform + +`ArrayLibrary` no longer supplies NumPy-style implementations. In 0.5.0 every +library inherited `self.array_module.linspace(...)`, `.ones`, `.outer`, +`.floor`, `.max(axis=...)`, `data.astype(...)` and so on, whether or not its +module had those functions. They now live in parents a library declares: + +| Parent | Supplies | Inherited by | +|---|---|---| +| `NumpyArrayApi` | every array operation, by NumPy's names on `array_module` | NumPy, CuPy, JAX | +| `ArrayMath` | declares `linspace`, `ones`, `outer`, `floor` (abstract) | `NumpyArrayApi`, torch, TensorFlow | +| `InPlaceAssignment` | `assign` by `array[index] = values` | `MutableArrays` (NumPy, CuPy), torch, pyclesperanto | + +`MutableArrays` is now a `NumpyArrayApi` and `InPlaceAssignment` library. +`accepts`, `max`, `min`, `mean`, `sum`, `assign`, `_stack`, `_cast`, +`_reshape`, `_broadcast_to`, `_ones_like` and `_logical_and` are abstract in +`ArrayLibrary`. Torch and TensorFlow state the few they used to inherit; JAX +drops four overrides that only repeated the NumPy spelling. + +What changes for callers: `PyclesperantoLibrary` has no `linspace`, `ones`, +`outer` or `floor`. In 0.5.0 `ArrayLibrary.named("pyclesperanto").linspace(...)` +existed and failed inside the call; it is now an `AttributeError` at the +attribute. Ask `isinstance(library, ArrayMath)`. Nothing else is renamed or +removed, and every other built-in library answers every operation as before. + +What changes for a library class written outside the package: it defines the +operations above or inherits `NumpyArrayApi` (`ADDING_NEW_FRAMEWORKS.md` shows +both). + +### A decorated call in main memory does not touch the accelerator + +The wrapper a library decorator puts around a function (`execution_wrapper`) +reads the device from the call's arguments. In 0.5.0 it asked the library for +its selected device whenever the machine had an accelerator, which for torch +calls `torch.cuda.current_device()`: CUDA was initialised and a stream opened +for a function working on host tensors. + +- Arguments held on accelerator N: the call runs on this thread's stream for + device N, with out-of-memory recovery for that device. (0.5.0 used the + stream of the selected device even when the data was on another one.) +- Arguments in main memory, or no array of a library that has main memory + (torch, TensorFlow, JAX): the function is called as written. Nothing of the + accelerator is asked, and there is no out-of-memory retry. +- A library without main memory (CuPy, pyclesperanto) called without one of + its arrays runs on its selected device, as before. + +### Typing + +`mypy` runs with `warn_return_any` and without `ignore_missing_imports`. Every +library the package imports is declared in the repository's `stubs/` +directory, which lists what arraybridge calls and comes before any installed +copy, so the result is the same on every machine. A library class states that +its `module` is that library; `NumpyArrayApi` names its module by the +`NumpyFunctions` Protocol. For a class written outside the package this means: +`NumpyArrayApi` and `MutableArrays` leave `array_module` abstract, +`AcceleratorDevices` asks for `_device_if_accepted(value)`, `CudaLibrary` for +`create_stream()`, and a `DLPackImport` library that cannot read a payload says +so in `_reads_dlpack(payload)` instead of returning `NotImplemented` from +`_import_dlpack`. + +Annotations that now say what was already returned: a device or stream scope +is an `AbstractContextManager[object]` (CuPy's device and stream are their own +context managers), `from_dlpack` returns the array or `NotImplementedType`, +and `ArrayLibrary.array_module` is an `object` in the base class. + +A value of the wrong kind is now refused with `TypeError` in four places: +`MemoryContractAttribute.read` when the stored value is neither a string nor +`None`, `TorchLibrary.stream_scope` for an object that is neither a +`torch.cuda.Stream` nor `None`, a dtype lookup in torch or TensorFlow whose +attribute is not a dtype of that library, and `DeviceMove.convert` called with +two different libraries. + +`MemoryType`'s metaclass overrides are annotated with `EnumMeta`'s signatures; +its `__call__` override is gone (`Enum` already asks `_missing_`). + +### Not verified on a GPU + +0.5.1 was tested with NumPy on Python 3.10, 3.12 and 3.13 and with the torch +CPU wheel on 3.12. CuPy, TensorFlow, JAX and pyclesperanto, and torch on CUDA, +were exercised only through the fakes in `tests/fakes.py`. `TESTING.md` has +the commands for a GPU machine. + ## 0.5.0 Each supported array library is now one class of one family, @@ -121,26 +205,3 @@ nominal type of an array of any supported library; `library.array_class` is an array of that library. Native array classes are registered with them when a library first meets one. -`mypy` runs with `warn_return_any` and without `ignore_missing_imports`. Every -library the package imports is declared in the repository's `stubs/` -directory, which lists what arraybridge calls and comes before any installed -copy, so the result is the same on every machine. A library class states that -its `module` is that library; `NumpyArrayApi` names its module by the -`NumpyFunctions` Protocol. For a class written outside the package this means: -`NumpyArrayApi` and `MutableArrays` leave `array_module` abstract, -`AcceleratorDevices` asks for `_device_if_accepted(value)`, `CudaLibrary` for -`create_stream()`, and a `DLPackImport` library that cannot read a payload says -so in `_reads_dlpack(payload)` instead of returning `NotImplemented` from -`_import_dlpack`. - -Annotations that now say what was already returned: a device or stream scope -is an `AbstractContextManager[object]` (CuPy's device and stream are their own -context managers), `from_dlpack` returns the array or `NotImplementedType`, -and `ArrayLibrary.array_module` is an `object` in the base class. - -A value of the wrong kind is now refused with `TypeError` in four places: -`MemoryContractAttribute.read` when the stored value is neither a string nor -`None`, `TorchLibrary.stream_scope` for an object that is neither a -`torch.cuda.Stream` nor `None`, a dtype lookup in torch or TensorFlow whose -attribute is not a dtype of that library, and `DeviceMove.convert` called with -two different libraries. diff --git a/TESTING.md b/TESTING.md index 318784b..83100ea 100644 --- a/TESTING.md +++ b/TESTING.md @@ -37,6 +37,41 @@ that declaration-owned import defaults work. Framework warnings about shared CUDA plugin registration should be recorded separately from value, device, or shape failures. +## Commands for a GPU machine + +0.5.0's GPU run predates three later fixes, and 0.5.1 was not run on a GPU. +On a machine with CUDA 12 and about 15 GB of free disk: + +```bash +git clone https://github.com/OpenHCSDev/ArrayBridge.git && cd ArrayBridge +git checkout v0.5.1 +uv venv -p 3.12 .venv-gpu +uv pip install -p .venv-gpu/bin/python -e ".[dev,gpu]" +.venv-gpu/bin/python -m pytest -q --no-cov -ra +.venv-gpu/bin/python -m pytest -q --no-cov -ra -m "gpu or cupy or torch or tensorflow or jax or pyclesperanto" +.venv-gpu/bin/python -m pytest -q --no-cov -ra tests/test_operation_soundness.py +.venv-gpu/bin/python host_call_check.py +``` + +`tests/test_operation_soundness.py` calls every array operation of every +installed library on its first device; no skip should name an installed +library. `host_call_check.py` is: + +```python +import torch +from arraybridge import torch as on_torch + +@on_torch +def add_one(image): + return image + 1 + +add_one(torch.ones(2, 3, 3)) +assert not torch.cuda.is_initialized(), "a host call initialised CUDA" +if torch.cuda.is_available(): + assert add_one(torch.ones(2, 3, 3, device="cuda:0")).is_cuda +print("ok") +``` + ## Interpreting CI The main CI workflow runs the test suite across its declared Python and OS diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index 28efb42..5376e43 100644 --- a/src/arraybridge/__init__.py +++ b/src/arraybridge/__init__.py @@ -10,7 +10,7 @@ NumPy/numcodecs import cost. """ -__version__ = "0.5.0" +__version__ = "0.5.1" _LAZY_EXPORTS: dict[str, str] = { "MemoryType": ".types",