stringjax_tools.vmap

Manual cached vmap/jit factories.

Cached wrappers snapshot static keyword configuration at construction time. For predictable cache reuse, prefer immutable configuration values.

vmapping_func(func[, in_axes])	Return a fresh jax.jit(jax.vmap(...)) wrapper around func.
vmapping_func_cached(func[, in_axes])	Return a cached jax.jit(jax.vmap(...)) wrapper around func.
clear_vmap_caches()	Clear wrappers cached by vmapping_func_cached().