Consider using numpyro.contrib.control_flow.scan over jax.lax.scan

[damonbayer]

Our use of jax.lax.scan has been fine so far, but the built-in numpyro.contrib.control_flow.scan offers additional features/compatibility with numpyro.

[dylanhmorris]

We have now mostly moved to this, but with the caveats to be explored in #444

[sbidari]

We use jax.lax.scan in the Infection classes, should these be switched to numpyro.contrib.control_flow.scan?

PyRenew/pyrenew/latent/infection_functions.py

Line 48 in 3d06ab9

latest, all_infections = jax.lax.scan(f=incidence_func, init=I0, xs=Rt)

PyRenew/pyrenew/latent/infection_functions.py

Line 184 in 3d06ab9

latest, infs_and_R_adj = jax.lax.scan(

[dylanhmorris]

I think for deterministic scans we should probably stick with jax.lax.scan but let's discuss in #444