stringjax_tools.jit.jit_with_static_args

jit_with_static_args(func, static_argnums=())

Wrap func with jax.jit and explicit positional static arguments.

Parameters:

• func (Callable[..., Any]) – Function to JIT-compile.

• static_argnums (tuple[int, ...]) – Positional argument indices to mark static.

Returns:

JIT-compiled wrapper around ``func``.

Return type:

Callable[..., Any]