stringjax_tools.auto_vectorise

Rank-checked automatic vectorisation helpers.

The central helper is auto_vmap(), a conservative decorator for functions whose single-sample input ranks are known. At call time, selected arguments are checked before JAX tracing. Rank r is treated as one sample, rank r + 1 as a paired leading-axis batch, and checked batched arguments must share the same leading size.

Core decorator

ArgSpec([shape])

Optional exact trailing sample-shape specification.

auto_vmap([sample_ranks, sample_shapes, ...])

Decorate a single-sample function with conservative automatic batching.

Defaults and cache management

set_auto_vmap_defaults(**sample_ranks)	Update global default sample ranks used by auto_vmap().
get_auto_vmap_defaults()	Return a copy of the active global sample-rank defaults.
reset_auto_vmap_defaults()	Clear all global sample-rank defaults.
auto_vmap_defaults(**sample_ranks)	Temporarily update global sample-rank defaults inside a with block.
set_auto_vmap_default_shapes(**sample_shapes)	Update global default exact sample-shape specifications.
get_auto_vmap_default_shapes()	Return a copy of the active global sample-shape defaults.
reset_auto_vmap_default_shapes()	Clear all global sample-shape defaults.
auto_vmap_default_shapes(**sample_shapes)	Temporarily update global sample-shape defaults inside a with block.
clear_auto_vmap_caches([clear_jax])	Clear cached wrappers created by auto_vmap().