When we use tfp.vi.fit_surrogate_posterior instead of our own (very simple) VI loss something gets mixed up and we end up with wrong results. When a singleton sample dimension is added this replicates the behaviour of fit_surrogate_posterior.
Just adding sample_size and importance_sample_size arguments (>1) doesn't cause any run time errors so the issue must be a silent mixup of axes. Does the fact that nucleotide frequencies come out correct but in the wrong order give us a clue what might be happening?
When we use
tfp.vi.fit_surrogate_posteriorinstead of our own (very simple) VI loss something gets mixed up and we end up with wrong results. When a singleton sample dimension is added this replicates the behaviour offit_surrogate_posterior.Just adding
sample_sizeandimportance_sample_sizearguments (>1) doesn't cause any run time errors so the issue must be a silent mixup of axes. Does the fact that nucleotide frequencies come out correct but in the wrong order give us a clue what might be happening?