How to use kfac to train two probabilistic models jointly?

[wangleiphy]

In my application, I need to jointly optimize two probabilistic models. They contribute to two different terms in the final loss function.

I am wondering what would be the recommended pattern of using kfac ?
More specifically, does it make sense to invoke kfac_jax.register_normal_predictive_distribution twice (for the two probabilistic models respectively) ?

Thanks in advance!

[wangleiphy]

Perhaps a more concrete example of my question:

How train a VAE using kfac-jax?

[botev]

VAEs are currently not supported, general speaking KFAC has been developed only in the context of supervised learning. Anything else would require some further research to know what is the appropriate way to use it.

[wangleiphy]

Thanks!