Manual vmap helpers

Use vmapping_func and vmapping_func_cached when the vectorisation pattern is already explicit. These helpers are the manual counterpart to auto_vmap: you provide in_axes, and the package returns a jax.jit(jax.vmap(...)) wrapper.

import jax.numpy as jnp
import numpy as np

from stringjax_tools.vmap import clear_vmap_caches, vmapping_func, vmapping_func_cached

Fresh and cached wrappers

vmapping_func creates a fresh wrapper every time. vmapping_func_cached reuses a wrapper for the same function, axes, and static keyword configuration.

def single_sample(x, y, scale=1.0):
    return scale * (x + y)

fresh = vmapping_func(single_sample, in_axes=(0, None), scale=2.0)
cached_1 = vmapping_func_cached(single_sample, in_axes=(0, None), scale=2.0)
cached_2 = vmapping_func_cached(single_sample, in_axes=(0, None), scale=2.0)

print('fresh result:', fresh(jnp.arange(3.0), 10.0))
print('cached result:', cached_1(jnp.arange(3.0), 10.0))
print('same cached object:', cached_1 is cached_2)

Static keyword configuration

Static keyword configuration is snapshotted when the cached wrapper is built. This keeps the cache key and closure behavior consistent even if a mutable object passed by the caller is changed later. In production code, immutable configuration values such as tuples are clearer.

def shifted(x, shifts):
    return x + shifts[0]

shifts = [3.0]
wrapped = vmapping_func_cached(shifted, in_axes=0, shifts=shifts)
shifts[0] = 100.0

print(wrapped(jnp.arange(3.0)))
clear_vmap_caches()