Add Bayesian Flows :rocket:
Bayesian Flows in zuko would be a very useful feature for uncertainty estimation with normalizing flows.
This PR introduces Bayesian Flow in zuko with minimal changes: `BayesianLinear` layers are provided as an alternative way of building the MLPs at the base of zuko transformations
All flows can be transformed to their Bayesian version by adding `bayesian=True` to the building arguments.
An utility function `total_KL_divergence(model)` is added in the `zuko.utils` to compute the total KL divergence from all the bayesian layers in a model to be added to the loss.
Let me know if you agree with the design of this feature! I added tests and a tutorial ;)
P.S.: CNF flows are not working for the moment but I haven't investigated why yet.
合并状态:未合并 31 条评论