Reproducibility issue of MRVI

[justjhong]

Moving issue reported by @Nusob888 from mrvi for tracking.

Dear Yosef lab members,

Thanks for MRVI, I have been trying it out finally.

Unfortunately I am experiencing some reproducibility issues. I first noticed this when I accidentally refreshed a Jupyter notebook and re-ran the code. The embeddings were very different.
Since then I have tried a numerous amount of methods to set seed and I cannot seem to produce anything consistent.

This would be a big issue for reproducibility if the u embeddings were ever to be used for clustering (which I had hoped to do) and I presume it would also effect how the abundance and differential expression analysis perform?

I have attached 4 htmls of my code run on the PBMC 3K dataset.

Any updates on this would be great to help me keep track of the tool, as I would very much like to use it, but have concerns around this issue.

SCVI_MRVI_reproducibility_scvi_1.pdf
SCVI_MRVI_reproducibility_mrvi_2.pdf
SCVI_MRVI_reproducibility_mrvi_1.pdf
SCVI_MRVI_reproducibility_scvi_2.pdf

[justjhong]

Found the solution. This only occurs on CUDA, not on CPU. I was able to reproduce the issue then solve it using the solution posted here: jax-ml/jax#13672.

os.environ['XLA_FLAGS'] = '--xla_gpu_deterministic_ops=true'

I also tested this after removing all the lines fixing seeds except for scvi.settings.seed, and this was sufficient with the flag to achieve determinism. Closing now, but feel free to open if the issue persists.

@Nusob888

[canergen]

Thanks @justjhong. Should we maybe provide the flag inside scvi-tools or add it to the tutorial?

[justjhong]

sure, let's add it into the seed function of the scviconfig. I think it slows down performance a little so it should only be applied when setting seed manually

[Nusob888]

Amazing! thank you both!