jaxvacua.util.auto_vmap

auto_vmap(sample_ranks=None, *, sample_shapes=None, validate_shapes=True, jit=True, **named_sample_ranks)

JAXVacua-local wrapper around stringjax_tools.auto_vmap().

The decorator applies JAXVacua’s canonical sample ranks and exact trailing shapes only to argument names present in the decorated function. Local declarations override the defaults, for example @auto_vmap(fluxes=1).

Return type:

Callable[[Callable[..., Any]], Callable[..., Any]]