jarnvid 0.2.1__tar.gz

jarnvid-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Noah Syrkis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jarnvid-0.2.1/PKG-INFO

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: jarnvid
+Version: 0.2.1
+Summary: Complex structures as frozen JAX dataclasses, made algebraic: write for one instance, vectorize over any batch.
+Keywords: jax,vmap,vectorize,jaxtyping,pytree,dataclass,batching
+Author: Noah Syrkis
+Author-email: Noah Syrkis <noah@syrkis.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Dist: jax>=0.9.0
+Requires-Dist: jaxtyping>=0.3.9
+Requires-Python: >=3.13
+Project-URL: Repository, https://github.com/syrkis/jarnvid
+Project-URL: Decisions, https://github.com/syrkis/jarnvid/tree/main/docs/adr
+Description-Content-Type: text/markdown
+
+# jarnvid
+
+Write a JAX function for one instance of a structure, and call it on any batch of them.
+`vectorize` reads which axes belong to one instance from the jaxtyping annotations you
+already write, and maps the function over the rest.
+
+```sh
+pip install jarnvid
+```
+
+Járnviðr is the iron wood of Norse myth. The structures it is for are complex things written
+as data (a force on a map, a doctrine, a molecule): frozen dataclasses of JAX arrays, nested in
+tuples and named sums.
+
+## Example
+
+```pycon
+>>> from dataclasses import dataclass
+>>> from typing import Annotated
+>>> import jax.numpy as jnp
+>>> from jax.tree_util import register_dataclass
+>>> from jaxtyping import Array, Float
+>>> from jarnvid import Extend, stack, take, vectorize
+>>> @register_dataclass
+... @dataclass(frozen=True)
+... class Unit:
+...     position: Float[Array, "*batch 2"]
+...     health: Float[Array, "*batch"]
+>>> type Force = Annotated[Unit, Extend("unit")]  # a side: its unit axis stays inside each call
+>>> @vectorize
+... def threat(force: Force, target: Float[Array, "*batch 2"]) -> Float[Array, "*batch"]:
+...     distance = jnp.linalg.norm(force.position - target, axis=-1)
+...     return (force.health / (1 + distance)).sum()
+>>> forces = Unit(position=jnp.zeros((5, 3, 2)), health=jnp.ones((5, 3)))  # 5 scenarios, 3 units each
+>>> threat(forces, jnp.ones((5, 2))).shape  # one target per scenario
+(5,)
+>>> threat(forces, jnp.ones(2)).shape  # one target for all: it reaches every call as is
+(5,)
+>>> threat(stack([forces, forces]), jnp.ones((2, 5, 2))).shape  # batch axes nest freely
+(2, 5)
+>>> take(forces, 0).position.shape  # the first scenario's force
+(3, 2)
+
+```
+
+## Why not `vmap`
+
+`jax.vmap` maps one axis, chosen per argument with `in_axes`. Several batch axes mean nested
+`vmap`s, and a dataclass whose fields carry different numbers of batch axes means spelling
+out an `in_axes` tree that mirrors it. `equinox.filter_vmap` chooses what to map by type, not
+by shape. jarnvid reads the shapes from the annotations instead:
+
+- One function signature gives any number of batch axes, broadcast **from the left** across
+  all arguments (docs/adr/0001).
+- A leaf with no batch axes (a map, a rule table) reaches every call once, never copied.
+- `Extend("unit")` keeps a named axis inside each call, so a function can see a whole side.
+- The same decorator maps in chunks (`Vectorizer(placement=Chunked(64))`, bounded memory) or
+  shards over a device mesh (`Vectorizer(placement=Sharded("data"))`, padding a batch that does
+  not divide evenly).
+
+The price is annotation discipline: every mapped array is annotated, starting with `*batch`.
+
+## Rules
+
+- Arguments are PyTrees: tuples, dicts (`dict[str, T]`, `Mapping[str, T]`), registered
+  dataclasses (fields marked `static` are left alone), named sums (the arm matching the value)
+  and type aliases. Arrays in any other container are an error.
+- What is mapped is read from the annotations alone. Plain Python values and unregistered
+  dataclasses (configs) reach every call as they are.
+- A leaf's batch axes beyond the ones shared by all arguments stay inside each call.
+
+Between calls, structures are reshaped by their batch axes alone (docs/adr/0002): `take`,
+`stack`, `concat`, `expand` (singleton axes, for outer products) and `batch_shape`. Leaves
+without batch axes stay once.
+
+| File | Concept |
+|---|---|
+| `annotation.py` | Reading annotations: which trailing axes of each leaf one call sees |
+| `vectorizer.py` | `vectorize` and `Vectorizer`: broadcast, flatten, map, reshape |
+| `placement.py` | Where the batch runs: `Local`, `Chunked`, or `Sharded` over a mesh axis |
+| `structure.py` | `take`, `stack`, `concat`, `expand`, `batch_shape` on batched structures |
+
+## Develop
+
+`just test` (the tests and this README's examples), `just sharded` (four forced CPU devices),
+`just check` (ruff and vulture), `just build` (the wheel and sdist). The scope is fixed: one
+decorator and the structure operations around it.
+
+## Cite
+
+See `CITATION.cff`. MIT licence.
+
+## Open
+
+- JAX 0.12: jarnvid passes unchanged. With `jax>=0.12`, add a law that `Sharded` adds no
+  collectives (`jax.sharding.no_collectives`), except for the result gather after padding.

jarnvid-0.2.1/README.md

@@ -0,0 +1,96 @@
+# jarnvid
+
+Write a JAX function for one instance of a structure, and call it on any batch of them.
+`vectorize` reads which axes belong to one instance from the jaxtyping annotations you
+already write, and maps the function over the rest.
+
+```sh
+pip install jarnvid
+```
+
+Járnviðr is the iron wood of Norse myth. The structures it is for are complex things written
+as data (a force on a map, a doctrine, a molecule): frozen dataclasses of JAX arrays, nested in
+tuples and named sums.
+
+## Example
+
+```pycon
+>>> from dataclasses import dataclass
+>>> from typing import Annotated
+>>> import jax.numpy as jnp
+>>> from jax.tree_util import register_dataclass
+>>> from jaxtyping import Array, Float
+>>> from jarnvid import Extend, stack, take, vectorize
+>>> @register_dataclass
+... @dataclass(frozen=True)
+... class Unit:
+...     position: Float[Array, "*batch 2"]
+...     health: Float[Array, "*batch"]
+>>> type Force = Annotated[Unit, Extend("unit")]  # a side: its unit axis stays inside each call
+>>> @vectorize
+... def threat(force: Force, target: Float[Array, "*batch 2"]) -> Float[Array, "*batch"]:
+...     distance = jnp.linalg.norm(force.position - target, axis=-1)
+...     return (force.health / (1 + distance)).sum()
+>>> forces = Unit(position=jnp.zeros((5, 3, 2)), health=jnp.ones((5, 3)))  # 5 scenarios, 3 units each
+>>> threat(forces, jnp.ones((5, 2))).shape  # one target per scenario
+(5,)
+>>> threat(forces, jnp.ones(2)).shape  # one target for all: it reaches every call as is
+(5,)
+>>> threat(stack([forces, forces]), jnp.ones((2, 5, 2))).shape  # batch axes nest freely
+(2, 5)
+>>> take(forces, 0).position.shape  # the first scenario's force
+(3, 2)
+
+```
+
+## Why not `vmap`
+
+`jax.vmap` maps one axis, chosen per argument with `in_axes`. Several batch axes mean nested
+`vmap`s, and a dataclass whose fields carry different numbers of batch axes means spelling
+out an `in_axes` tree that mirrors it. `equinox.filter_vmap` chooses what to map by type, not
+by shape. jarnvid reads the shapes from the annotations instead:
+
+- One function signature gives any number of batch axes, broadcast **from the left** across
+  all arguments (docs/adr/0001).
+- A leaf with no batch axes (a map, a rule table) reaches every call once, never copied.
+- `Extend("unit")` keeps a named axis inside each call, so a function can see a whole side.
+- The same decorator maps in chunks (`Vectorizer(placement=Chunked(64))`, bounded memory) or
+  shards over a device mesh (`Vectorizer(placement=Sharded("data"))`, padding a batch that does
+  not divide evenly).
+
+The price is annotation discipline: every mapped array is annotated, starting with `*batch`.
+
+## Rules
+
+- Arguments are PyTrees: tuples, dicts (`dict[str, T]`, `Mapping[str, T]`), registered
+  dataclasses (fields marked `static` are left alone), named sums (the arm matching the value)
+  and type aliases. Arrays in any other container are an error.
+- What is mapped is read from the annotations alone. Plain Python values and unregistered
+  dataclasses (configs) reach every call as they are.
+- A leaf's batch axes beyond the ones shared by all arguments stay inside each call.
+
+Between calls, structures are reshaped by their batch axes alone (docs/adr/0002): `take`,
+`stack`, `concat`, `expand` (singleton axes, for outer products) and `batch_shape`. Leaves
+without batch axes stay once.
+
+| File | Concept |
+|---|---|
+| `annotation.py` | Reading annotations: which trailing axes of each leaf one call sees |
+| `vectorizer.py` | `vectorize` and `Vectorizer`: broadcast, flatten, map, reshape |
+| `placement.py` | Where the batch runs: `Local`, `Chunked`, or `Sharded` over a mesh axis |
+| `structure.py` | `take`, `stack`, `concat`, `expand`, `batch_shape` on batched structures |
+
+## Develop
+
+`just test` (the tests and this README's examples), `just sharded` (four forced CPU devices),
+`just check` (ruff and vulture), `just build` (the wheel and sdist). The scope is fixed: one
+decorator and the structure operations around it.
+
+## Cite
+
+See `CITATION.cff`. MIT licence.
+
+## Open
+
+- JAX 0.12: jarnvid passes unchanged. With `jax>=0.12`, add a law that `Sharded` adds no
+  collectives (`jax.sharding.no_collectives`), except for the result gather after padding.

jarnvid-0.2.1/pyproject.toml

@@ -0,0 +1,59 @@
+[project]
+name = "jarnvid"
+version = "0.2.1"
+description = "Complex structures as frozen JAX dataclasses, made algebraic: write for one instance, vectorize over any batch."
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.13"
+keywords = [
+    "jax",
+    "vmap",
+    "vectorize",
+    "jaxtyping",
+    "pytree",
+    "dataclass",
+    "batching",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering",
+    "Typing :: Typed",
+]
+dependencies = [
+    "jax>=0.9.0",
+    "jaxtyping>=0.3.9",
+]
+
+[[project.authors]]
+name = "Noah Syrkis"
+email = "noah@syrkis.com"
+
+[project.urls]
+Repository = "https://github.com/syrkis/jarnvid"
+Decisions = "https://github.com/syrkis/jarnvid/tree/main/docs/adr"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-xdist",
+    "hypothesis",
+    "ruff>=0.17.4",
+    "vulture>=2.16",
+]
+
+[build-system]
+requires = ["uv_build>=0.13,<0.14"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 128
+
+[tool.pytest.ini_options]
+testpaths = [
+    "tests",
+    "README.md",
+]
+addopts = "-n auto --doctest-glob=README.md"

jarnvid-0.2.1/pyproject.toml.orig

@@ -0,0 +1,46 @@
+[project]
+name = "jarnvid"
+version = "0.2.1"
+description = "Complex structures as frozen JAX dataclasses, made algebraic: write for one instance, vectorize over any batch."
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Noah Syrkis", email = "noah@syrkis.com" }]
+requires-python = ">=3.13"
+keywords = ["jax", "vmap", "vectorize", "jaxtyping", "pytree", "dataclass", "batching"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering",
+    "Typing :: Typed",
+]
+dependencies = [
+    "jax>=0.9.0",
+    "jaxtyping>=0.3.9",
+]
+
+[project.urls]
+Repository = "https://github.com/syrkis/jarnvid"
+Decisions = "https://github.com/syrkis/jarnvid/tree/main/docs/adr"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-xdist",
+    "hypothesis",
+    "ruff>=0.17.4",
+    "vulture>=2.16",
+]
+
+[build-system]
+requires = ["uv_build>=0.13,<0.14"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 128
+
+[tool.pytest.ini_options]
+# The README's examples run as doctests, so it cannot fall behind the code.
+testpaths = ["tests", "README.md"]
+addopts = "-n auto --doctest-glob=README.md"

jarnvid-0.2.1/src/jarnvid/__init__.py

@@ -0,0 +1,15 @@
+"""Járnviðr, the iron wood: complex structures as data, algebraic in JAX.
+
+Doctrine, forces, molecules: frozen dataclasses of arrays. `vectorize` reads a function's
+jaxtyping annotations (`*batch` first), broadcasts the leading batch axes of its PyTree
+arguments, and maps the function over them; `Extend` keeps named semantic axes inside each call;
+`Vectorizer` maps the batch in chunks, or shards it over a mesh axis. `take`, `stack`, `concat` and `expand` reshape structures
+along their batch axes only.
+"""
+
+from .annotation import Extend
+from .placement import Chunked, Local, Sharded
+from .structure import batch_shape, concat, expand, stack, take
+from .vectorizer import Vectorizer, vectorize
+
+__all__ = ["Chunked", "Extend", "Local", "Sharded", "Vectorizer", "batch_shape", "concat", "expand", "stack", "take", "vectorize"]

jarnvid-0.2.1/src/jarnvid/annotation.py

@@ -0,0 +1,174 @@
+"""Reading annotations: for each leaf of an argument, how many trailing axes belong to one call.
+
+An argument's annotation is walked together with its value, once: through type aliases,
+`Annotated[..., Extend(...)]`, the one arm of a union that admits the value, tuples, dicts and
+dataclasses, down to jaxtyping arrays annotated `*batch ...`. The result is a PyTree of the
+value's structure whose leaves are `Core` (the array's last `ndim` axes belong to each call; the
+rest, at least one, are batch axes) or `Fixed` (passed to every call unchanged).
+"""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, fields, is_dataclass, replace
+from types import UnionType
+from typing import Annotated, TypeAliasType, TypeIs, Union, get_args, get_origin, get_type_hints
+
+from jax import tree
+from jaxtyping import AbstractArray, PyTree
+
+# Any runtime annotation (a class, alias, union or Annotated form); PEP 747's TypeForm, which
+# Python 3.13 lacks, is the precise type.
+type TypeForm = object
+
+
+@dataclass(frozen=True, init=False)
+class Extend:
+    """Insert semantic axes immediately after every nested ``*batch`` axis."""
+
+    axes: tuple[str, ...]
+
+    def __init__(self, *axes: str) -> None:
+        if not axes or any(not axis for axis in axes):
+            raise ValueError("Extend requires one or more nonempty axis names")
+        object.__setattr__(self, "axes", axes)
+
+
+@dataclass(frozen=True)
+class Core:
+    """An array leaf whose last `ndim` axes belong to each call; the rest are batch axes."""
+
+    ndim: int
+
+
+@dataclass(frozen=True)
+class Fixed:
+    """A leaf passed to every call unchanged."""
+
+
+type Rank = Core | Fixed
+
+
+def ranks(value: PyTree, annotation: TypeForm, path: str, extended: tuple[str, ...] = ()) -> PyTree:
+    """`value`'s structure with a `Rank` at every leaf, checking each array against its annotation."""
+    annotation, axes = resolve(value, annotation, path)
+    extended += axes
+    if value is None:
+        return None
+    if is_array(annotation):
+        return array_rank(value, annotation, path, extended)
+    if isinstance(value, tuple):
+        items = zip(value, tuple_annotations(annotation, len(value), path), strict=True)
+        return tuple(ranks(item, inner, f"{path}[{index}]", extended) for index, (item, inner) in enumerate(items))
+    if isinstance(value, dict):
+        inner = mapping_annotation(annotation, path)
+        return {key: ranks(item, inner, f"{path}[{key!r}]", extended) for key, item in value.items()}
+    if is_dataclass(value):
+        return dataclass_ranks(value, annotation, path, extended)
+    if hasattr(value, "shape"):
+        raise TypeError(f"{path} must use a jaxtyping annotation beginning with '*batch'")
+    if any(hasattr(leaf, "shape") for leaf in tree.leaves(value)):
+        raise TypeError(f"{path} holds arrays in a {type(value).__name__}, which is not traversed: use a tuple, dict or dataclass")
+    return tree.map(lambda _: Fixed(), value)
+
+
+def array_rank(value: PyTree, annotation: type[AbstractArray], path: str, extended: tuple[str, ...]) -> Rank:
+    """An array's `Rank`: `Core` if it has batch axes, else `Fixed`."""
+    if not hasattr(value, "shape"):
+        raise TypeError(f"{path} is annotated as an array but received {type(value).__name__}")
+    core = core_ndim(annotation, path) + len(extended)
+    if value.ndim < core:
+        raise ValueError(f"{path} has shape {value.shape}, which is smaller than {extended_dims(annotation, extended)!r}")
+    # A leaf without batch axes reaches every call as is, never copied across the batch.
+    return Core(core) if value.ndim > core else Fixed()
+
+
+def dataclass_ranks(value: PyTree, annotation: TypeForm, path: str, extended: tuple[str, ...]) -> PyTree:
+    """A registered dataclass with the `Rank`s of its non-static fields; an unregistered one (a
+    config) is `Fixed`, unless it holds arrays."""
+    if any(leaf is value for leaf in tree.leaves(value)):
+        if any(hasattr(leaf, "shape") for leaf in tree.leaves([getattr(value, field.name) for field in fields(value)])):
+            raise TypeError(f"{path}: register {type(value).__name__} as a PyTree (jax.tree_util.register_dataclass)")
+        return Fixed()
+    if not isinstance(annotation, type) or not is_dataclass(annotation):
+        raise TypeError(f"{path} must have an annotated dataclass type")
+    hints = get_type_hints(annotation, include_extras=True)
+    mapped = (field.name for field in fields(value) if not field.metadata.get("static", False))
+    return replace(value, **{name: ranks(getattr(value, name), hints.get(name), f"{path}.{name}", extended) for name in mapped})
+
+
+def mapping_annotation(annotation: TypeForm, path: str) -> TypeForm:
+    """The value type of a ``dict[K, V]`` or ``Mapping[K, V]`` annotation."""
+    if get_origin(annotation) not in (dict, Mapping) or len(get_args(annotation)) != 2:
+        raise TypeError(f"{path} must have an annotated dict[K, V] or Mapping[K, V] type")
+    return get_args(annotation)[1]
+
+
+def resolve(value: PyTree, annotation: TypeForm, path: str) -> tuple[TypeForm, tuple[str, ...]]:
+    """Unwrap aliases and ``Extend`` metadata, taking a union's one arm that admits ``value``."""
+    annotation, axes = extension(annotation)
+    while get_origin(annotation) in (Union, UnionType):
+        arms = tuple(arm for arm in get_args(annotation) if admits(extension(arm)[0], value))
+        if len(arms) != 1:
+            raise TypeError(f"{path} ({type(value).__name__}) matches {len(arms)} arms of {annotation}, not one")
+        annotation, arm_axes = extension(arms[0])
+        axes += arm_axes
+    return annotation, axes
+
+
+def extension(annotation: TypeForm) -> tuple[TypeForm, tuple[str, ...]]:
+    """Separate a runtime annotation from its (possibly nested) semantic-axis metadata."""
+    axes: tuple[str, ...] = ()
+    seen: set[int] = set()
+    while True:
+        if isinstance(annotation, TypeAliasType):
+            if id(annotation) in seen:
+                raise TypeError(f"Cyclic type alias: {annotation.__name__}")
+            seen.add(id(annotation))
+            if annotation.__type_params__:
+                raise TypeError(f"Generic type aliases are not supported: {annotation.__name__}")
+            annotation = annotation.__value__
+        elif isinstance(get_origin(annotation), TypeAliasType):
+            raise TypeError(f"Generic type aliases are not supported: {annotation}")
+        elif get_origin(annotation) is Annotated:
+            annotation, *metadata = get_args(annotation)
+            axes += tuple(axis for item in metadata if isinstance(item, Extend) for axis in item.axes)
+        else:
+            return annotation, axes
+
+
+def admits(annotation: TypeForm, value: PyTree) -> bool:
+    """Whether one union arm describes a runtime value: arrays by shape, others by class."""
+    if is_array(annotation):
+        return hasattr(value, "shape")
+    origin = get_origin(annotation) or annotation
+    return isinstance(origin, type) and isinstance(value, origin)
+
+
+def tuple_annotations(annotation: TypeForm, length: int, path: str) -> tuple[TypeForm, ...]:
+    """Expand a fixed or variadic tuple annotation over its runtime values."""
+    if get_origin(annotation) is not tuple:
+        raise TypeError(f"{path} must have an annotated tuple type")
+    annotations = get_args(annotation)
+    if len(annotations) == 2 and annotations[1] is Ellipsis:
+        return (annotations[0],) * length
+    if len(annotations) != length:
+        raise TypeError(f"{path} has {length} values but {len(annotations)} annotated tuple fields")
+    return annotations
+
+
+def is_array(annotation: TypeForm) -> TypeIs[type[AbstractArray]]:
+    """Whether an annotation is a concrete jaxtyping array type."""
+    return isinstance(annotation, type) and issubclass(annotation, AbstractArray)
+
+
+def core_ndim(annotation: type[AbstractArray], path: str) -> int:
+    """The axes to the right of the required leading ``*batch``."""
+    dimensions = annotation.dim_str.split()
+    if annotation.index_variadic != 0 or not dimensions or dimensions[0] != "*batch":
+        raise TypeError(f"{path} must use a jaxtyping annotation beginning with '*batch'")
+    return len(annotation.dims) - 1
+
+
+def extended_dims(annotation: type[AbstractArray], extended: tuple[str, ...]) -> str:
+    """The effective annotation, semantic axes included, for error messages."""
+    batch, *core = annotation.dim_str.split()
+    return " ".join((batch, *extended, *core))

jarnvid-0.2.1/src/jarnvid/placement.py

@@ -0,0 +1,93 @@
+"""Where the flattened batch runs: on the local devices, or split over one mesh axis."""
+
+from dataclasses import dataclass
+from typing import Protocol, assert_never
+
+import jax
+import jax.numpy as jnp
+from jax import tree
+from jax.sharding import NamedSharding, PartitionSpec
+from jaxtyping import PyTree
+
+from .annotation import Core
+
+
+@dataclass(frozen=True, slots=True)
+class Local:
+    """Map the flattened batch on the current devices."""
+
+
+@dataclass(frozen=True, slots=True)
+class Chunked:
+    """Map the flattened batch `batch` instances at a time (``lax.map``): the same result as
+    ``Local`` in bounded memory."""
+
+    batch: int
+
+
+@dataclass(frozen=True, slots=True)
+class Sharded:
+    """Split the one flattened batch axis over the active mesh's ``axis``."""
+
+    axis: str
+
+
+type Placement = Local | Chunked | Sharded
+
+
+class Apply[R](Protocol):
+    """One call of the wrapped function, on the mapped arguments by name."""
+
+    def __call__(self, values: dict[str, PyTree], /) -> R: ...
+
+
+def run[R](placement: Placement, apply: Apply[R], values: dict[str, PyTree], ranks: dict[str, PyTree], size: int) -> R:
+    """`apply` mapped over the leading axis (of length `size`) of every `Core` leaf of `values`."""
+    axes = tree.map(lambda rank: 0 if isinstance(rank, Core) else None, ranks)
+    mapped = jax.vmap(apply, in_axes=(axes,))
+    if size == 0:
+        return mapped(values)
+    match placement:
+        case Local():
+            return mapped(values)
+        case Chunked(batch=batch):
+
+            def at(index: jax.Array) -> R:
+                return apply(tree.map(lambda leaf, rank: leaf[index] if isinstance(rank, Core) else leaf, values, ranks))
+
+            return jax.lax.map(at, jnp.arange(size), batch_size=batch)
+        case Sharded(axis=axis):
+            return sharded(axis, mapped, values, ranks, size)
+        case _:
+            assert_never(placement)
+
+
+def sharded[R](axis: str, mapped: Apply[R], values: dict[str, PyTree], ranks: dict[str, PyTree], size: int) -> R:
+    """`mapped` over the batch split across the active mesh's `axis`.
+
+    The batch is padded to a multiple of the axis with copies of its last instance, and the
+    padding cut from a replicated result (a sharded axis cannot be cut unevenly). Unbatched
+    arrays are replicated: shard_map cannot close over arrays sharded on Explicit axes.
+    """
+    mesh = jax.sharding.get_abstract_mesh()
+    if axis not in mesh.shape:
+        raise ValueError(f"vectorize over Sharded({axis!r}) requires an active mesh with that axis")
+    padding = -size % mesh.shape[axis]
+
+    def spec(leaf: PyTree, rank: PyTree) -> PartitionSpec | None:
+        if not hasattr(leaf, "shape"):
+            return None
+        return PartitionSpec(axis) if isinstance(rank, Core) else PartitionSpec()
+
+    def place(leaf: PyTree, rank: PyTree) -> PyTree:
+        if not hasattr(leaf, "shape"):
+            return leaf
+        if isinstance(rank, Core):
+            leaf = jnp.pad(leaf, [(0, padding)] + [(0, 0)] * (leaf.ndim - 1), mode="edge")
+        return jax.sharding.reshard(leaf, NamedSharding(mesh, spec(leaf, rank)))
+
+    local = jax.shard_map(mapped, mesh=mesh, in_specs=(tree.map(spec, values, ranks),), out_specs=PartitionSpec(axis), check_vma=False)
+    result = local(tree.map(place, values, ranks))
+    if not padding:
+        return result
+    return tree.map(lambda leaf: jax.sharding.reshard(leaf, PartitionSpec())[:size], result)

jarnvid-0.2.1/src/jarnvid/structure.py

@@ -0,0 +1,95 @@
+"""Batch-aware operations on structures: `take`, `stack`, `concat`, `expand`, `batch_shape`.
+
+A structure is a registered dataclass whose annotations say which axes of each leaf are batch
+axes, as `vectorize` reads them. Only leaves with batch axes move: a leaf without them (a terrain
+kept once) stays once, and plain values (Python scalars) must agree.
+"""
+
+from collections.abc import Sequence
+from typing import Protocol
+
+import jax.numpy as jnp
+from jax import Array, tree
+from jax.typing import ArrayLike
+from jaxtyping import PyTree
+
+from .annotation import Core, Fixed, Rank, ranks
+from .vectorizer import batch_shapes, shared
+
+
+def layout(value: PyTree) -> PyTree:
+    """`value`'s `Rank`s, read from its own class's annotations."""
+    return ranks(value, type(value), type(value).__name__)
+
+
+def batch_shape(value: PyTree) -> tuple[int, ...]:
+    """The batch `vectorize` maps `value` over."""
+    return shared(batch_shapes(type(value).__name__, value, layout(value)))
+
+
+def take[T](value: T, index: ArrayLike) -> T:
+    """The instance (or instances) at `index` along the leading batch axis."""
+    size = leading(value)
+    return tree.map(lambda leaf, rank: lead(leaf, rank, size)[index] if isinstance(rank, Core) else leaf, value, layout(value))
+
+
+def stack[T](values: Sequence[T]) -> T:
+    """`values` along a new leading batch axis."""
+    return merge(values, lambda leaves, _: jnp.stack(jnp.broadcast_arrays(*leaves)))
+
+
+def concat[T](values: Sequence[T]) -> T:
+    """`values` joined along their leading batch axis."""
+    sizes = tuple(map(leading, values))
+    return merge(values, lambda leaves, ranks: jnp.concatenate([lead(*parts) for parts in zip(leaves, ranks, sizes, strict=True)]))
+
+
+def expand[T](value: T, before: int = 0, after: int = 0) -> T:
+    """`value` with singleton batch axes before and after its batch axes; an outer product is
+    `f(expand(xs, after=1), expand(ys, before=1))`."""
+    rank = len(batch_shape(value))
+
+    def insert(leaf: PyTree, r: PyTree) -> PyTree:
+        if not isinstance(r, Core):
+            return leaf
+        return leaf.reshape((1,) * before + leaf.shape[:rank] + (1,) * after + leaf.shape[rank:])
+
+    return tree.map(insert, value, layout(value))
+
+
+class Join(Protocol):
+    """Combine one leaf from each structure into one leaf."""
+
+    def __call__(self, leaves: tuple[Array, ...], ranks: tuple[Rank, ...], /) -> Array: ...
+
+
+def merge[T](values: Sequence[T], join: Join) -> T:
+    """Leaf by leaf: agreeing plain values, and one leaf without batch axes that every structure
+    shares, stay once; the rest are `join`ed."""
+    count = len(values)
+
+    def leaf(*items: PyTree) -> PyTree:
+        leaves, rank = items[:count], items[count:]
+        if not all(hasattr(item, "shape") for item in leaves):
+            if any(item != leaves[0] for item in leaves):
+                raise ValueError(f"Plain values differ across the structures: {leaves}")
+            return leaves[0]
+        if all(isinstance(r, Fixed) and item is leaves[0] for item, r in zip(leaves, rank, strict=True)):
+            return leaves[0]
+        return join(leaves, rank)
+
+    return tree.map(leaf, *values, *map(layout, values))
+
+
+def leading(value: PyTree) -> int:
+    """The size of `value`'s leading batch axis."""
+    shape = batch_shape(value)
+    if not shape:
+        raise ValueError(f"{type(value).__name__} has no batch axes")
+    return shape[0]
+
+
+def lead(leaf: Array, rank: Rank, size: int) -> Array:
+    """`leaf` with a leading batch axis of `size`, broadcast if it has one, added if not."""
+    shape = leaf.shape[1:] if isinstance(rank, Core) else leaf.shape
+    return jnp.broadcast_to(leaf, (size, *shape))

jarnvid-0.2.1/src/jarnvid/vectorizer.py

@@ -0,0 +1,96 @@
+"""`vectorize`: write a function for one instance, call it on any batch.
+
+Each argument's annotation gives, for every array leaf, the trailing axes one call sees; the
+leading axes left over are batch axes. Batch axes broadcast across leaves **aligned from the
+left**: the shared batch is the broadcast of every leaf's first `r` batch axes, `r` the fewest
+any leaf has, and a leaf's further batch axes stay inside each call (docs/adr/0001). The batch
+is flattened to one axis, mapped (`placement`), and reshaped back.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from functools import wraps
+from inspect import signature
+from math import prod
+from typing import get_type_hints
+
+import jax.numpy as jnp
+from jax import Array, tree
+from jax.tree_util import keystr
+from jaxtyping import PyTree
+
+from .annotation import Core, ranks
+from .placement import Local, Placement, run
+
+# A function whose signature vectorize preserves.
+type Function[**P, R] = Callable[P, R]
+
+
+def vectorize[**P, R](fn: Function[P, R]) -> Function[P, R]:
+    """Vectorize a PyTree function from its jaxtyping annotations."""
+    return Vectorizer()(fn)
+
+
+@dataclass(frozen=True, slots=True)
+class Vectorizer:
+    """A configured ``vectorize``: where the batch runs. What is mapped is read from the
+    annotations alone; `Extend` keeps named axes inside each call."""
+
+    placement: Placement = field(default_factory=Local)
+
+    def __call__[**P, R](self, fn: Function[P, R]) -> Function[P, R]:
+        call = signature(fn)
+        hints = get_type_hints(fn, include_extras=True)
+
+        @wraps(fn)
+        def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+            bound = call.bind(*args, **kwargs)
+            bound.apply_defaults()
+            arguments = bound.arguments
+            layout = {name: ranks(arguments[name], hints.get(name), name) for name in arguments}
+            batches = {name: batch_shapes(name, arguments[name], rank) for name, rank in layout.items()}
+            batch = shared({path: shape for shapes in batches.values() for path, shape in shapes.items()})
+            if not batch:
+                return fn(*args, **kwargs)
+            values = {name: tree.map(lambda leaf, r: flatten(leaf, r, batch), arguments[name], rank) for name, rank in layout.items()}
+
+            def apply(inner: dict[str, PyTree]) -> R:
+                one = call.bind_partial()
+                one.arguments.update(arguments | inner)
+                return fn(*one.args, **one.kwargs)
+
+            result = run(self.placement, apply, values, layout, prod(batch))
+            return tree.map(lambda leaf: leaf.reshape(batch + leaf.shape[1:]), result)
+
+        return wrapped
+
+
+def batch_shapes(name: str, value: PyTree, rank: PyTree) -> dict[str, tuple[int, ...]]:
+    """The batch shape (axes before the core) of every batched array leaf, by path."""
+    pairs = zip(tree.leaves_with_path(value), tree.leaves(rank), strict=True)
+    return {name + keystr(path): leaf.shape[: leaf.ndim - r.ndim] for (path, leaf), r in pairs if isinstance(r, Core)}
+
+
+def shared(shapes: dict[str, tuple[int, ...]]) -> tuple[int, ...]:
+    """The broadcast of every batched leaf's first `r` batch axes, `r` the fewest any has."""
+    if not shapes:
+        return ()
+    rank = min(map(len, shapes.values()))
+    try:
+        return jnp.broadcast_shapes(*(shape[:rank] for shape in shapes.values()))
+    except ValueError:
+        leaves = ", ".join(f"{path} {shape}" for path, shape in shapes.items())
+        raise ValueError(
+            f"Batch axes do not broadcast in their first {rank}, aligned from the left: {leaves}. "
+            "An axis meant to stay inside each call belongs in the annotation, or in Extend."
+        ) from None
+
+
+def flatten(leaf: PyTree, rank: PyTree, batch: tuple[int, ...]) -> PyTree:
+    """A `Core` leaf broadcast to the shared batch, which is flattened to one leading axis; a
+    leaf's batch axes beyond the shared ones stay inside the call."""
+    if not isinstance(rank, Core):
+        return jnp.asarray(leaf) if hasattr(leaf, "shape") else leaf
+    array: Array = leaf
+    rest = array.shape[len(batch) : array.ndim - rank.ndim] + array.shape[array.ndim - rank.ndim :]
+    return jnp.broadcast_to(array, batch + rest).reshape((prod(batch), *rest))