AnyTensor

Portable tensor ops across NumPy, JAX, PyTorch, and TensorFlow, with a focus on segment / GNN primitives.

Write a helper once; run it on whatever tensor the caller already has. Ordinary math uses the Python Array API via array-api-compat. Segment reductions stay on thin input-adaptive backends.

Docs (motivation, GAT-style case study, design, API): https://swamidass.github.io/anytensor/ — or uv run --group docs mkdocs serve from a checkout (docs/).

We follow Semantic Versioning: breaking changes require a major bump. Portability is backed by cross-backend / symbolic fuzz, a coverage gate, and pytest-run docs examples (Design).

Install

pip install "anytensor @ git+https://github.com/swamidass/anytensor.git"
# optional backends (NumPy is a core dependency)
pip install "anytensor[jax]" "anytensor[torch]" "anytensor[tensorflow]"
# or
pip install "anytensor[all]"

Requires Python ≥3.10. Backend floors: NumPy ≥1.24, JAX ≥0.4.32, PyTorch ≥2.1, TensorFlow ≥2.13.

Quick start

import anytensor as at
import numpy as np

x = np.arange(12.0).reshape(3, 4)
seg_ids = np.array([0, 0, 1])
y = at.segment_sum(x, seg_ids, num_segments=2)

The same call works on JAX / Torch / TF tensors. num_segments is required (JAX convention).

anytensor.jraph is a portable jraph: GraphsTuple, batching/padding, and GraphNetwork on any backend. Nested feature trees use anytensor.tree (jax.tree API; pure Python, NumPy is the only binary dep). See the docs Jraph and Tree sections.

Read next:

• Home / motivation — why AnyTensor, GAT neighbor-softmax case study across four backends
• Jraph — portable GraphsTuple / GraphNetwork
• Tree — nest helpers (pure Python + NumPy) for graphs and any structured record
• Design — principles, edge cases, testing as contract, SemVer
• Usage — promotion, segment helpers, torch.compile, typing
• Surprising differences — NaN / ±inf / graph / GPU gotchas from fuzz
• API reference — generated from docstrings
• Contributing — tests, fuzz, docs build

Docs / tests (checkout)

uv sync --extra all --group dev --group docs
uv run mkdocs serve
uv run pytest -m "not fuzz" --cov=anytensor --cov-report=term-missing

License

MIT. Backend dispatch patterns adapted from einops; see NOTICE.