Metadata-Version: 2.4
Name: ndlike
Version: 0.1.2
Summary: A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch.
Author-email: Clemens Damke <clemens@cortys.de>
License-Expression: MIT
Project-URL: source, https://github.com/Cortys/ndlike
Project-URL: tracker, https://github.com/Cortys/ndlike/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: 3.15
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=2.1.0
Requires-Dist: flextype>=0.3.0
Provides-Extra: torch
Requires-Dist: torch>=2.6.0; extra == "torch"
Provides-Extra: jax
Requires-Dist: jax>=0.8.1; extra == "jax"
Dynamic: license-file

# ndlike

NumPy arrays, JAX arrays, and PyTorch tensors provide similar primitives for
indexing, reshaping, stacking, and numerical operations. `ndlike` provides tools
to build wrappers around these types that behave like the underlying arrays,
but whose elements represent custom composite values—such as vectors or
distributions—instead of individual numbers.

Protocols and reusable base classes handle array interfaces and backend function
dispatch. Protected axes distinguish collection dimensions from each element's
internal dimensions, so indexing, reshaping, or stacking a collection preserves
the structure of its elements.

## Installation

Requires Python 3.12 or newer.

```bash
pip install ndlike
pip install 'ndlike[torch]'
pip install 'ndlike[jax]'
```

NumPy and flextype are core dependencies. Torch and JAX are optional and are
never imported by the backend-independent modules. Import backend modules explicitly.

## Array protocols and base classes

`ndlike.array_like.ArrayLike` describes indexing, shape, and metadata rather than
just convertibility to an array (as in `numpy.typing.ArrayLike`). It does not
require the full numerical Array API. `ArrayLikeWithoutSize` covers the shared
interface without NumPy/JAX's integer `size` property, since Torch uses a method.

Backend-specific APIs build on these protocols:

| Module | Protocol | Conversion helper | Conversion hook |
| --- | --- | --- | --- |
| `ndlike.array_like` | `NumpyLike` | `to_numpy_like` | `__array_like__` |
| `ndlike.torch_like` | `TorchLike` | `to_torch_like` | `__torch_like__` |
| `ndlike.jax_like` | `JaxLike` | `to_jax_like` | `__jax_like__` |

Each module also provides matching `Implementation` and `Convertible` classes,
such as `NumpyLikeImplementation` and `NumpyLikeConvertible`.

Use the array protocols to type inputs, or subclass an implementation base to
reuse convenience methods while supplying your representation's metadata and
operation-dispatch hooks. NumPy uses `__array_function__` / `__array_ufunc__`,
Torch uses `__torch_function__`, and ndlike supplies `__jax_function__` for JAX.

Conversion helpers call the conversion hook when an input satisfies the matching
convertible protocol: the core array interface plus the hook, not the hook alone.
This lets custom types preserve their structure instead of becoming raw arrays.
Otherwise, the helpers fall back to `numpy.asanyarray`, `torch.as_tensor`, and
`jax.numpy.asarray`, respectively. NumPy retains `__array_like__` for backward
compatibility and consistency with NumPy's `__array__`.

## Protected axes

Protected-axis base classes give structured objects array-like behavior over
their batch dimensions while preserving trailing internal dimensions. For
example, a batch of six 3D vectors should have logical shape `(6,)`, not `(6, 3)`:

```python
from dataclasses import dataclass
from typing import ClassVar

import numpy as np

from ndlike.protected_axis.numpy import NumpyAxisProtected


@dataclass
class Vectors(NumpyAxisProtected[np.ndarray]):
    array: np.ndarray
    protected_axes: ClassVar[dict[str, int]] = {"array": 1}


vectors = Vectors(np.ones((6, 3)))
assert vectors.shape == (6,)
assert vectors.protected_shape == (3,)
assert vectors[0].array.shape == (3,)

grid = vectors.reshape(2, 3)
assert grid.shape == (2, 3)
assert grid.array.shape == (2, 3, 3)

stacked = np.stack([vectors, vectors])
assert isinstance(stacked, Vectors)
assert stacked.shape == (2, 6)
```

`protected_axes` maps field names to the number of trailing axes to preserve.
Multiple fields can have different internal shapes while sharing the same batch
shape; indexing and supported shape operations update them together. Dataclasses
work out of the box; other classes can override `with_protected_values` to
reconstruct results.

Numerical operations are opt-in: use `permitted_functions` (for example,
`{np.mean}`) and, for NumPy ufuncs, `permitted_ufuncs` (for example,
`{np.add: ["__call__", "reduce"]}`). Supported reductions operate on batch axes,
leaving internal axes intact.

The corresponding backend classes are `TorchAxisProtected` in
`ndlike.protected_axis.torch` and `JaxAxisProtected` in
`ndlike.protected_axis.jax`.

## JAX function dispatch

JAX has no native equivalent of NumPy's `__array_function__`. For custom objects,
use wrappers from `ndlike.jax_functions`, such as `jax_reshape`, `jax_stack`, and
`jax_mean`, instead of the corresponding `jax.numpy` functions. They dispatch to
`__jax_function__` on custom types and call JAX normally for native arrays.
`jax_function_dispatch` lets you add the same dispatch behavior to your own
functions. `JaxLikeImplementation` also provides dataclass-based pytree support
for JAX transformations such as `jax.jit` and `jax.vmap`.

## Development

```bash
uv sync --all-extras
uv run pytest
uv run prek install
uv run prek run --all-files
uv build
```

Dependency groups follow flextype: `lint` contains Ruff and ty, `test` contains
pytest, and `dev` includes both groups plus prek. Development Torch wheels use
the CPU index on Linux and Windows. Type checking requires both backend extras.

`scripts/check_backend_naming.py` enforces backend prefixes on definitions
(BKN001) and restricts runtime Torch/JAX imports to their respective
backend-prefixed modules (BKN002). Type-only imports are exempt. Intentional
exceptions use an explicit `# noqa: BKN001` or `# noqa: BKN002` on the opening line.
