JIT helpers

This notebook covers the small static-argument helpers. They are intentionally thin wrappers around jax.jit; use them when they make a call boundary clearer than spelling out jax.jit(..., static_argnums=...) directly.

import jax.numpy as jnp
import numpy as np

from stringjax_tools.jit import is_static, jit_with_dynamic_static_args, jit_with_static_args

Explicit static arguments

Use jit_with_static_args when the static positional arguments are known in advance. This is the production-friendly helper because the static argument pattern is fixed.

def polynomial(x, offset, power):
    return (x + offset) ** power

compiled = jit_with_static_args(polynomial, static_argnums=(2,))
print(compiled(jnp.array([1.0, 2.0]), 1.0, 2))

Dynamic static arguments

jit_with_dynamic_static_args infers static positional arguments at call time. That is useful while exploring, but changing static positions can cause repeated JAX compilation, so explicit static arguments are preferable in stable code.

dynamic = jit_with_dynamic_static_args(polynomial)
print(dynamic(jnp.array([1.0, 2.0]), 1.0, 2))

print('python scalar static:', is_static(1.0))
print('numpy array static:', is_static(np.ones(2)))
print('jax array static:', is_static(jnp.ones(2)))