Automatic vectorisation

This notebook introduces auto_vmap, the conservative automatic vectorisation helper. The intended pattern is to write one clear single-sample function and then declare the sample ranks of the arguments that may carry a leading batch axis.

The helper implements paired batching only. If several checked arguments are batched, their leading dimensions must agree; unbatched checked arguments are broadcast with in_axes=None.

import jax.numpy as jnp
import numpy as np

from stringjax_tools.auto_vectorise import auto_vmap

Rank-only batching

The declarations below mean that moduli and fluxes are vector-valued single samples, while tau is scalar-valued. A shape such as (N, h) for moduli is therefore interpreted as a leading-axis batch of N samples.

@auto_vmap(moduli=1, tau=0, fluxes=1, jit=False)
def observable(moduli, tau, fluxes, scale=1.0):
    return scale * (jnp.sum(moduli) + tau + jnp.sum(fluxes))

single = observable(jnp.array([1.0, 2.0]), jnp.array(3.0), jnp.arange(4.0))
paired = observable(jnp.ones((3, 2)), jnp.arange(3.0), jnp.ones((3, 4)))
broadcast = observable(jnp.ones((3, 2)), jnp.array(3.0), jnp.ones(4))

print('single:', single)
print('paired:', paired)
print('broadcast:', broadcast)
print('axes:', observable.auto_vmap_axes(jnp.ones((3, 2)), jnp.array(3.0), jnp.ones(4)))

Exact shape validation

Ranks decide whether an argument is batched. Exact trailing dimensions are an optional diagnostic layer. String dimensions such as "h12" are resolved from self or another bound object; derived dimensions should be callables.

class ToyModel:
    def __init__(self, h12, n_fluxes):
        self.h12 = h12
        self.n_fluxes = n_fluxes

    @auto_vmap(
        sample_ranks={'moduli': 1, 'fluxes': 1},
        sample_shapes={
            'moduli': 'h12',
            'fluxes': lambda bound: (2 * bound['self'].n_fluxes,),
        },
        validate_shapes=True,
        jit=False,
    )
    def diagnostic(self, moduli, fluxes):
        return jnp.sum(jnp.abs(moduli) ** 2) + 0.01 * jnp.sum(fluxes**2)

model = ToyModel(h12=3, n_fluxes=2)
values = model.diagnostic(jnp.ones((4, 3)), jnp.ones((4, 4)))
print(values)

Shape errors are raised before JAX traces the function, which makes them easier to interpret than downstream tracing errors.

try:
    model.diagnostic(jnp.ones((4, 2)), jnp.ones((4, 4)))
except ValueError as exc:
    print(type(exc).__name__ + ':', exc)