stringjax_tools.vmap.vmapping_func_cached

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

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

The cache key includes the callable, in_axes, and a robust key for bound keyword configuration. Keyword values may be unhashable; they are used by the closure, while their cache-key representation is produced by stringjax_tools._utils.cache_key_value().

Return type:

Callable[..., Any]