Refactor fft_conv_general_dilated to match jax.lax.conv_general_dilated signature

[ASEM000]

Currently, fft based convolution function looks like this

def fft_conv_general_dilated(
    x: jnp.ndarray,
    w: jnp.ndarray,
    strides: tuple[int, ...],
    padding: tuple[tuple[int, int], ...],
    groups: int,
    dilation: tuple[int, ...],
) -> jnp.ndarray:

An improvement would be to match jax.lax.conv_general_dilated signature; this requires adding the lhs_dilation option and renaming the function arguments.

Related issues and motivation for FFT implementation :

google/jax#6343
google/jax#5227
google/jax#7284 -> serket.nn.FFTFilter2D can improve on this
google-deepmind/dm_pix#46 serket.nn.FFTFilter2D can improve on this for large kernel size

[chrisflesher]

Was wondering if you have any plans to integrate fft_conv_general_dilated with the jax repo? It would be nice to speed up jax.lax.conv_general_dilated for CPU's...

[ASEM000]

Sounds reasonable, I will reopen the issue when time permits.

[chrisflesher]

Cool, thanks!