stringjax_tools.jit

Small JIT helpers for functions with Python-side static arguments.

jit_with_static_args(func[, static_argnums])	Wrap func with jax.jit and explicit positional static arguments.
jit_with_dynamic_static_args(func)	JIT func while deciding static positional arguments at call time.
is_static(arg)	Heuristic test for whether arg should be treated as JAX-static.