Add boundary='wrap' support for jax.scipy.signal.convolve2d and jax.scipy.signal.correlate2d

[rajasekharporeddy]

Currently jax.scipy.signal.convolve2d and jax.scipy.signal.correlate2d functions support only boundary='fill'. This PR will add boundary='wrap' support for both the functions.

Fixes #7276

[jakevdp]

Looks great! The last thing we need is to squash the changes into a single commit (see https://jax.readthedocs.io/en/latest/contributing.html#single-change-commits-and-pull-requests). Let me know if you need help with that process

[rajasekharporeddy]

Squashed all the commits into a single one.

[jakevdp]

Looks great, thanks!

[jakevdp]

We're seeing errors when running a larger number of test cases. You can reproduce them locally using e.g.

JAX_NUM_GENERATED_CASES=50 pytest -n auto tests/scipy_signal_test.py -k testConvolutions

It looks like it's an issue with boundary='wrap' combined with mode='full'.

[jakevdp]

Example failure:

testConvolutions2DNotValidMode7 (mode='full', boundary='wrap', op='convolve2d', dtype=<class 'numpy.float64'>, xshape=(3, 4), yshape=(2, 5)) (shard 21)content_copy
---
  File "jax/tests/scipy_signal_test.py", line 163, in testConvolutions2DNotValidMode
    self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=False,
  File "jax/_src/public_test_util.py", line 127, in _assert_numpy_allclose
    np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
  File "numpy/testing/_private/utils.py", line 1504, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "numpy/testing/_private/utils.py", line 797, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=1e-12, atol=1e-12

Mismatched elements: 4 / 32 (12.5%)
Max absolute difference: 11.06572735
Max relative difference: 2.87831292
 x: array([[ -7.221209, -14.248762, -25.984588, -41.801186,  -7.221209,
        -14.248762, -25.984588, -41.801186],
       [-47.816622,  -3.279326, -16.912216,  42.092961, -47.816622,...
 y: array([[  3.844519, -14.248762, -25.984588, -41.801186,  -7.221209,
        -14.248762, -25.984588, -41.801186],
       [-45.311671,  -3.279326, -16.912216,  42.092961, -47.816622,...

[rajasekharporeddy]

It is fixed now. It is the problem with padding in1. It appears that the pad_width should have same value in all dimensions which is the maximum shape value among the input shapes.

[jakevdp]

That's surprising to me – say you are convolving two (1000, 2) arrays. Padding it out to (2000, 1002) seems unnecessary, because we really only need 4 elements in the second dimension.

Can we fix this without blowing up the memory requirements in cases like this?

[rajasekharporeddy]

I will check and let you know, if it is possible or not.

[rajasekharporeddy]

I could not find any other way to fix the issue without padding it with maximum value in all dimensions.

[jakevdp]

I would consider this a blocker to merging this PR: we need to figure out how to implement this without excessive padding.