Implement gaussian_filter1d and gaussian_filter

[adonath]

This PR implements jax.scipy.ndimage.gaussian_filter1d and jax.scipy.ndimage.gaussian_filter, addresses #7284. The implementation is mostly along the lines of what is done in scipy, whoever with some smaller differences:

• No support for the output argument
• No support for boundary modes other than "constant", "wrap" and "reflect", however they become available once they are supported by jnp.pad
• Support of a method option to choose between "direct" and "fft" based convolutions.

Here you can find a notebook doing some comparisons of the accuracy and performance:
https://github.com/adonath/notebooks-public/blob/master/jax-vs-scipy-gaussian-filter.ipynb

The takeaways are:

• The direct mode is asymptotically ~3x times slower than Scipy
• The FFT is the same performance for the 1d case as Scipy (would need to check why this make sense...)
• The ND-FFT option is asymptotically faster than Scipy, which makes sense, as Scipy relies on direct convolution (correlate1d I think...)

ToDo:

• Implement tests
• Investigate difference in the "reflect" mode

[jakevdp]

Thanks! However, in https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html#scipy-ndimage we determined that scipy.ndimage is out-of-scope for JAX. I wonder if this functionality might fit in dm-pix or another JAX-adjacent library?

[adonath]

Thanks @jakevdp! This makes sense and the JEP is very clear on the rational behind. I guess It would make sense to close related issues such as #7284 or #3647 and with a reference to the JEP just to avoid other contributors hopping on it...

I would be happy to contribute this to dm-pix, however it seems the scope of dm-pix is really batched images with only a single channel dimension (typical for deep learning), and not the general n-dimensional approach from scipy. Not sure where the right place is though...