Optimization reproducibility #2
Description
Running SCATrEx twice with the same random seed yields different results. This does not happen when we do not do any ELBO optimization, and just change the parameters in the move proposals and do ELBO evaluations with those. It is irrespective of local in the run_search being True or False. For some reason, some moves lead to ELBO optimization steps with small divergences between runs, and this accumulates leading to different trees in the end.
Commenting out the ancestor nodes' contribution from the node weights seems to fix it. So it might be a numerical instability problem that XLA does not deal with deterministically?