wcosmo.utils.maybe_jit

wcosmo.utils.maybe_jit(func, *jit_args, **jit_kwargs)

A decorator to jit the function if using jax.

This also allows arbitrary arguments to be passed through, e.g., to specify static arguments.

This function is pretty useful and so might make it into gwpopulation regardless of cosmology.