Refactoring SteinVI

[OlaRonning]

SVGD has been extended to a few varieties (ASVGD, GSVGD, HSVGD) which I want to include in the einstein module. This PR consists of the changes to SteinVI to allow for the extensions and introduction of SVGD and ASVGD. GSVGD and HSVGD will be in subsequent PRs.

To this end, I've:

1. Added a specialized constructor for SVGD.
   • The SVGD constructor does not require a guide because it's always a delta
   • It correctly sets the scaling on the attractive force to 1/m for m particles.
   • It doesn't allow users to change num_elbo_particles because it's always 1 for SVGD.
2. Added a setup_run method to SteinVI.
   • The method encapsulates a step of SteinVI, which inheriting constructors can manipulate.
3. Added a constructor for ASVGD.
   • ASVGD introduces an annealing schedule on the attractive force.
   • This inherits from SVGD and overwrites setup_run to change the loss_temperature.
4. Removed Jacobian projection from the Stein force
   • The projection is unnecessary as we do not allow AutoIAFNormal, AutoBNAFNormal, AutoDAIS, AutoSemiDAIS and AutoSurrogateLikelihoodDAIS.
   • This simplifies the force computation to attractive+repulsive.

Misc changes

1. The normalization factor in the ProbabilityProductKernel has been removed. The kernel is still a proper kernel; however, this version avoids vanishing/exploding when the guide variances deviate from 1 for "high" dimensional models.
2. Added a rng_key to the kernel. This is convenient when experimenting with kernels.
3. Removed the enum parameter as it is currently unsupported.
4. Kernel smoothing of the attractive force and the repulsion is now computed on particles in unconstraint space, consistent with particles moving in unconstraint space. NB: Constraint particles are removed from the test.
5. Added the manual computation for kernel tests as comments.
   • I inadvertently changed the value of one of the particles when doing the manual computations, which is why the evaluation has changed.

TODO

• Add docstrings to ASVGD and SVGD.
• Add manual computation to the IMQKernel and GraphicalKernel tests.
• Fix SteinVI documentation I'll do this in a separate PR.

[OlaRonning]

@fehiepsi Let me know if you want me to split it into multiple PRs.

[fehiepsi]

LGTM pending minor comments.

[fehiepsi]

This is deprecated. You can use jax.tree.flatten and jax.tree.map now.

[fehiepsi]

same as above

[fehiepsi]

same as above

[OlaRonning]

Thanks! I'll change it to jax.tree instead.