Metadata-Version: 2.4
Name: arraybridge
Version: 0.5.0
Summary: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion
Project-URL: Homepage, https://github.com/OpenHCSDev/arraybridge
Project-URL: Documentation, https://arraybridge.readthedocs.io
Project-URL: Repository, https://github.com/OpenHCSDev/arraybridge
Project-URL: Issues, https://github.com/OpenHCSDev/arraybridge/issues
Author-email: Tristan Simas <tristan.simas@mail.mcgill.ca>
License: MIT
License-File: LICENSE
Keywords: array,conversion,cupy,gpu,jax,numpy,pytorch,tensor,tensorflow
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Requires-Dist: metaclass-registry>=0.4.0
Requires-Dist: numpy>=1.20
Requires-Dist: python-introspect>=0.1.11
Provides-Extra: all
Requires-Dist: cupy>=10.0; extra == 'all'
Requires-Dist: jax>=0.3; extra == 'all'
Requires-Dist: jaxlib>=0.3; extra == 'all'
Requires-Dist: pyclesperanto>=0.10; extra == 'all'
Requires-Dist: tensorflow>=2.8; extra == 'all'
Requires-Dist: torch>=1.10; extra == 'all'
Provides-Extra: cupy
Requires-Dist: cupy>=10.0; extra == 'cupy'
Provides-Extra: dev
Requires-Dist: black>=23.0; extra == 'dev'
Requires-Dist: build>=1.0; extra == 'dev'
Requires-Dist: dill>=0.3.8; extra == 'dev'
Requires-Dist: mypy>=1.0; extra == 'dev'
Requires-Dist: packaging>=23.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Requires-Dist: twine>=5.0; extra == 'dev'
Provides-Extra: docs
Requires-Dist: sphinx-autodoc-typehints>=1.25; extra == 'docs'
Requires-Dist: sphinx-rtd-theme>=2.0; extra == 'docs'
Requires-Dist: sphinx>=7.0; extra == 'docs'
Provides-Extra: gpu
Requires-Dist: cucim-cu12<26.0.0,>=25.6.0; extra == 'gpu'
Requires-Dist: cupy-cuda12x<14.0.0,>=13.3.0; extra == 'gpu'
Requires-Dist: jax-cuda12-pjrt<0.6.0,>=0.5.3; extra == 'gpu'
Requires-Dist: jax-cuda12-plugin<0.6.0,>=0.5.3; extra == 'gpu'
Requires-Dist: jax<0.6.0,>=0.5.3; extra == 'gpu'
Requires-Dist: jaxlib<0.6.0,>=0.5.3; extra == 'gpu'
Requires-Dist: pyclesperanto>=0.17.1; extra == 'gpu'
Requires-Dist: tensorflow-probability[tf]<0.26.0,>=0.25.0; extra == 'gpu'
Requires-Dist: tensorflow<2.20.0,>=2.19.0; extra == 'gpu'
Requires-Dist: torch<2.8.0,>=2.6.0; extra == 'gpu'
Requires-Dist: torchvision<0.23.0,>=0.21.0; extra == 'gpu'
Provides-Extra: jax
Requires-Dist: jax>=0.3; extra == 'jax'
Requires-Dist: jaxlib>=0.3; extra == 'jax'
Provides-Extra: pyclesperanto
Requires-Dist: pyclesperanto>=0.10; extra == 'pyclesperanto'
Provides-Extra: tensorflow
Requires-Dist: tensorflow>=2.8; extra == 'tensorflow'
Provides-Extra: torch
Requires-Dist: torch>=1.10; extra == 'torch'
Description-Content-Type: text/markdown

# arraybridge

ArrayBridge converts arrays between NumPy, CuPy, PyTorch, TensorFlow, JAX and
pyclesperanto, and gives code one way to ask each library about devices,
memory and array operations.

Core dependencies are NumPy and metaclass-registry. The other libraries are
optional and are imported only when used.

Version 0.5.0 removed names without aliases. [CHANGELOG.md](CHANGELOG.md)
lists each removed name and what to call instead.

## Quick start

```python
import numpy as np

from arraybridge import HOST, convert_memory, detect_memory_type

value = np.arange(6).reshape(2, 3)
assert detect_memory_type(value) == "numpy"

result = convert_memory(value, source_type="numpy", target_type="numpy", device=HOST)
assert result is value
```

`convert_memory(data, source_type, target_type, device)` takes all four
arguments. `device` is a device of the target library: `HOST` for main memory,
`AcceleratorDevice(0)` for the library's first accelerator. A conversion to a
device the target library does not have fails; nothing is placed somewhere
else instead.

## Devices

A device is a value: `HOST` (main memory), `AcceleratorDevice(index)` (an
accelerator as its library numbers it) or `NamedDevice(name)` (a device a
library reports that is not one of its numbered accelerators, such as
pyclesperanto on a CPU OpenCL device). Indices belong to one library; the same
index in two libraries need not be the same hardware.

| Target library | Devices it accepts |
|---|---|
| NumPy | `HOST` |
| CuPy, pyclesperanto | `AcceleratorDevice(i)` |
| PyTorch, JAX, TensorFlow | `HOST` or `AcceleratorDevice(i)` |

Each library answers what it has on this machine:

```python
import numpy as np

from arraybridge import HOST, AcceleratorDevice, ArrayLibrary

numpy_library = ArrayLibrary.named("numpy")
assert numpy_library.host_devices() == (HOST,)
assert numpy_library.accelerator_devices() == ()
assert numpy_library.available_devices() == (HOST,)
assert numpy_library.device_of(np.zeros(3)) == HOST

assert AcceleratorDevice(0).label == "accelerator:0"
```

## Decorators

```python
from arraybridge import numpy

@numpy
def normalize(image):
    return image / max(float(image.max()), 1.0)

assert normalize.input_memory_type == "numpy"
```

The decorator named after a library sets `input_memory_type`,
`output_memory_type` and `execution_memory_type` on the function, adds the
`slice_by_slice` and `dtype_config` keyword arguments, and wraps the call with
what that library's devices need (per-thread CUDA streams for CuPy and
PyTorch, out-of-memory retries for libraries with accelerators). It does not
convert inputs or outputs between libraries and does not choose a device. The
caller converts with `convert_memory`.

## Stack utilities

```python
import numpy as np

from arraybridge import HOST, stack_slices, unstack_slices

slices = [np.zeros((8, 8)), np.ones((8, 8))]
stack = stack_slices(slices, memory_type="numpy", device=HOST)
restored = unstack_slices(stack, memory_type="numpy", device=HOST)
assert stack.shape == (2, 8, 8) and len(restored) == 2
```

`stack_slices` takes a non-empty list of 2D arrays. `unstack_slices` takes a
3D array. Both raise `ValueError` for a wrong shape or a device the target
library does not have.

## Which route a conversion takes

A conversion tries three mechanisms in order: `"device move"` (same library,
other device), `"dlpack"` (the target wraps the source's buffer) and
`"host copy"` (download to NumPy, then upload). `ConversionRoute.planned`
tells which one will be tried first. It reads only the library classes and the
device values, so it works without the libraries or the hardware:

```python
from arraybridge import HOST, AcceleratorDevice, ArrayLibrary, ConversionRoute

torch = ArrayLibrary.named("torch")
cupy = ArrayLibrary.named("cupy")
gpu = AcceleratorDevice(0)

assert ConversionRoute.planned(torch, gpu, cupy, gpu).mechanism == "dlpack"
assert ConversionRoute.planned(torch, HOST, cupy, gpu).mechanism == "host copy"
assert ConversionRoute.planned(torch, HOST, torch, gpu).mechanism == "device move"
```

To see what conversions actually ran, install a `ConversionCounters` observer.
It keeps a count, bytes and seconds per source, target and mechanism. This
example needs JAX:

```python
import numpy as np

from arraybridge import HOST, ConversionCounters, convert_memory

counters = ConversionCounters()
with counters.installed():
    convert_memory(np.zeros((4, 4), dtype=np.float32), "numpy", "jax", HOST)

row = counters.rows()[0]
assert (row["source_library"], row["target_library"]) == ("numpy", "jax")
assert (row["mechanism"], row["count"], row["nbytes"]) == ("host copy", 1, 64)
print(counters.format_table())
```

A call that returns its input unchanged is not counted. With no observer
installed nothing is timed.

## Array operations in the array's own library

`ArrayLibrary.named("torch")`, `ArrayLibrary.of(array)` and
`MemoryType.TORCH.library` all return the library object. Its methods run in
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`.

```python
import numpy as np

from arraybridge import ArrayLibrary


def max_projection(stack):
    library = ArrayLibrary.of(stack)
    return library.prepend_axis(library.max(stack, 0))


assert max_projection(np.ones((4, 8, 8))).shape == (1, 8, 8)
```

`library.module` is the library's module, imported on first use. Each library
keeps its own numeric rules: PyTorch, JAX and TensorFlow take means in
float32, and pyclesperanto reduces only along the leading axis.

`reshape`, `broadcast_to` and `ones_like` do not download pixels to NumPy.
pyclesperanto cannot change shape or broadcast natively, so those calls raise
`NotImplementedError` unless the shape is unchanged; its `ones_like` supports
one to three dimensions and its float32 and integer dtypes, not bool.

## Adding a library

A library is one class. See
[ADDING_NEW_FRAMEWORKS.md](ADDING_NEW_FRAMEWORKS.md).

## Installation

```bash
pip install arraybridge
pip install "arraybridge[torch]"
pip install "arraybridge[cupy]"
```

Documentation: <https://arraybridge.readthedocs.io>
