Vectorized and differentiable implementation of non-local means by Buades et. al for image denoising in JAX for EE367.
Approximately two orders of magnitude speed up with JIT compilation + GPU acceleration compared to a raw NumPy implementation.
python nlm.py
