adv_jax_math
