Add pupport forward-mode differentiation with autoguide

[juanitorduz]

Closes #2068

---

Summary

Adds forward_mode_differentiation parameter to all autoguide classes, fixing #2068.

AutoGuide._setup_prototype calls initialize_model to inspect the model and find valid initial parameters. Previously, it always used reverse-mode differentiation (the default), which fails for models containing JAX primitives that don't support reverse-mode AD (e.g., jax.lax.while_loop, jax.numpy.linalg.eigh). Users could pass forward_mode_differentiation=True to SVI.run/SVI.update to control differentiation of the ELBO loss, but the autoguide's internal initialize_model call was unaffected.

This PR adds a forward_mode_differentiation keyword argument (default False) to AutoGuide.__init__ and threads it through to all subclasses:

• AutoGuide (base class) — stores the flag and passes it to initialize_model in _setup_prototype
• AutoGuideList
• AutoNormal
• AutoDelta
• AutoDAIS
• AutoSurrogateLikelihoodDAIS — also passes the flag to the second initialize_model call for the surrogate model
• AutoSemiDAIS

AutoBatchedMixin uses *args, **kwargs and forwards automatically.

Usage

guide = AutoNormal(model, forward_mode_differentiation=True)
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
svi.run(rng_key, num_steps, forward_mode_differentiation=True)

[fehiepsi]

Thanks @juanitorduz!