PyTorch Implementation

[jh27kim]

Hi, thanks for sharing the code publicly.

Do you have plans to release the code in PyTorch? even just for the core implementations? (PosteriorDenoiser)

[francois-rozet]

Hi @jh27kim, thank you for your interest!

We do not plan to release the code of the experiments in PyTorch. However the moment matching posterior sampling (MMPS) method is available in the Azula library, which is based on PyTorch. The implementation of MMPS in Azula supports the GMRES solver, which we found to work better than the conjugate gradient method when the denoiser is not optimal. If you need help applying the method, please open a discussion over there.

MMPS was also benchmarked against other posterior sampling methods in https://github.com/francois-rozet/mmps-benchmark, but we do not recommend using this code.