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 a2aad93..3911db1 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. @@ -139,12 +176,11 @@ 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 | -| `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: @@ -154,18 +190,63 @@ 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. -- `_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` | `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). 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 @@ -181,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..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, @@ -120,3 +204,4 @@ 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. + diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 80455c2..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 @@ -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 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/TESTING.md b/TESTING.md index 752c3b6..83100ea 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 ``` @@ -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/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 d53a490..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). -warn_return_any = false +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/__init__.py b/src/arraybridge/__init__.py index 99d79d9..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", @@ -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/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 2f16fe2..432edff 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. @@ -15,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. """ @@ -40,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 @@ -51,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__) @@ -77,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) @@ -102,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: @@ -123,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 @@ -296,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 @@ -346,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) @@ -367,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) @@ -377,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: @@ -428,9 +531,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, @@ -552,18 +653,18 @@ 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 ) # ----- 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``.""" @@ -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): @@ -625,45 +728,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 +763,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] @@ -722,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) @@ -741,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 @@ -761,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) @@ -787,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: @@ -803,19 +897,55 @@ 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.""" + + if self.loaded_module() is None: + return None + for value in arguments: + 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.""" - return None, nullcontext() + 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[object]: + """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) - device, scope = self._thread_execution() + elif self.import_if_installed() is None or not self.accelerator_devices(): + return func(*args, **kwargs) + 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: @@ -833,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 @@ -871,14 +1000,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[object]: 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: @@ -898,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: @@ -919,9 +1055,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: @@ -972,34 +1106,127 @@ 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[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: NumpyNativeT, axis: Axis) -> NumpyNativeT: + return self.array_module.max(data, axis=axis) + + def min(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: + return self.array_module.min(data, axis=axis) + + def mean(self, data: NumpyNativeT, axis: Axis) -> NumpyNativeT: + return self.array_module.mean(data, axis=axis) + + 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) -> NumpyNativeT: + return self.array_module.linspace(start, stop, num, endpoint=endpoint) + + 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: NumpyNativeT, right: NumpyNativeT) -> NumpyNativeT: + return self.array_module.outer(left, right) + + def floor(self, value: float) -> NumpyNativeT: + return self.array_module.floor(value) + + def _stack(self, values: Sequence[NumpyNativeT]) -> NumpyNativeT: + return self.array_module.stack(tuple(values), axis=0) + + def _cast(self, data: NumpyNativeT, dtype: DType) -> NumpyNativeT: + return data.astype(dtype, copy=False) + + def _logical_and(self, left: NumpyNativeT, right: NumpyNativeT) -> NumpyNativeT: + return self.array_module.logical_and(left, right) + + def _reshape(self, data: NumpyNativeT, shape: tuple[int, ...]) -> NumpyNativeT: + return data.reshape(shape) + + def _broadcast_to(self, data: NumpyNativeT, shape: tuple[int, ...]) -> NumpyNativeT: + return self.array_module.broadcast_to(data, shape) + + 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[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: @@ -1021,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 @@ -1056,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: @@ -1079,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() @@ -1109,20 +1358,27 @@ class TorchLibrary( CudaLibrary["torch.Tensor"], DLPackImport["torch.Tensor"], DLPackProtocolExport["torch.Tensor"], + ArrayMath["torch.Tensor"], + InPlaceAssignment["torch.Tensor"], ): name = "torch" display_name = "PyTorch" 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: @@ -1135,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 ) @@ -1146,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: @@ -1184,10 +1450,15 @@ 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, self.module.dtype)) + + 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)) @@ -1196,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() @@ -1208,31 +1480,44 @@ 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: + 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) 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( HostDevices["tensorflow.Tensor"], AcceleratorDevices["tensorflow.Tensor"], DLPackImport["tensorflow.Tensor"], + ArrayMath["tensorflow.Tensor"], ): name = "tensorflow" display_name = "TensorFlow" @@ -1240,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( @@ -1276,24 +1567,21 @@ 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: 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) @@ -1305,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) @@ -1314,25 +1604,28 @@ 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) - 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] + def ones(self, shape: int | tuple[int, ...], dtype: DType) -> tensorflow.Tensor: + 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) @@ -1344,35 +1637,49 @@ 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) + + 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: - 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( @@ -1380,6 +1687,7 @@ class JaxLibrary( AcceleratorDevices["jax.Array"], DLPackImport["jax.Array"], DLPackProtocolExport["jax.Array"], + NumpyArrayApi["jax.Array"], ): name = "jax" display_name = "JAX" @@ -1387,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: @@ -1408,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: @@ -1429,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 @@ -1454,64 +1769,58 @@ 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) + 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 _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) + 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) -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") + if TYPE_CHECKING: + module = pyclesperanto + _native_dtype_names: ClassVar[frozenset[str]] = frozenset( {"float32", "int8", "int16", "int32", "uint8", "uint16", "uint32"} ) @@ -1542,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()) @@ -1569,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: @@ -1613,59 +1925,57 @@ 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: - module, np = self.module, _numpy() + def _scale_dtype(self, result: pyclesperanto.Array, target_dtype: DType) -> pyclesperanto.Array: + 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: @@ -1684,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/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/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 e29e0ab..553ea33 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,7 +164,17 @@ 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): + 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( @@ -171,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 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 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: