We could make the LatentVariableLayer interface more explicit by changing encoders from the status-quo "return parameters for the approximate posterior distribution" to a DistributionLambda that returns the posterior-approximating distribution itself.
This would make the behavior more explicit, and remove the need to keep track of what class the prior was (instead simply relying on tfp's kl_divergence implementations).
We could make the LatentVariableLayer interface more explicit by changing encoders from the status-quo "return parameters for the approximate posterior distribution" to a DistributionLambda that returns the posterior-approximating distribution itself.
This would make the behavior more explicit, and remove the need to keep track of what class the prior was (instead simply relying on tfp's kl_divergence implementations).