stringjax_tools.vmap.vmapping_func

vmapping_func(func, in_axes=None, **kwargs)

Return a fresh jax.jit(jax.vmap(...)) wrapper around func.

kwargs are bound inside the mapped closure. Each call constructs a new wrapper object, so repeated use should normally prefer vmapping_func_cached().

Return type:

Callable[..., Any]