Forward routines need to be updated to include residual terms from parameter expectations. This doesn't affect inference over latent variables, but does affect ELBO calculation for time series models. So the goal is forward and backward functions of the form:
output_dist, Res = transform.forward(input_dist, Res)
output_dist, Res = transform.backward(input_dist, Res)
where Res is updated to include terms not in output_dist that contribute to the ELBO calculation.
Forward routines need to be updated to include residual terms from parameter expectations. This doesn't affect inference over latent variables, but does affect ELBO calculation for time series models. So the goal is forward and backward functions of the form:
output_dist, Res = transform.forward(input_dist, Res)
output_dist, Res = transform.backward(input_dist, Res)
where Res is updated to include terms not in output_dist that contribute to the ELBO calculation.