stringjax_tools.auto_vectorise.auto_vmap

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

Decorate a single-sample function with conservative automatic batching.

Sample ranks decide batching. For an argument declared with sample rank r, rank r is treated as one unbatched sample and rank r + 1 is treated as a leading-axis batch. If several selected arguments are batched, their leading batch sizes must agree. Exact trailing-shape validation is optional and separate from batching detection.

Return type:

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