Documentation for convolve

[yoavram]

Hi, The doctoring for autograd.signal.convolve doesn't say much:

Wraps a function so that its gradient can be specified and its invocation
can be recorded. For examples, see the docs.

probably because it gives the doctoring of the decorator wraps, but anyway there is no doctoring in the code, either. This is unfortunate as the arguments are not "standard", in as much as they do not follow scipy.signal.convolve which it replaces or spicy.ndimage.convolve. We are left with reading the source.

[lina2002]

Parameter axes specifies axes on which convolution is being performed. It's a tuple consisting of a list of dimensions of the first matrix and a list of dimensions of the second matrix. Let's imagine e.g. that we have a list of 4 images and we want to apply the same convolution matrix to all of them:

>>> from autograd.scipy import signal
>>> import autograd.numpy as np
>>> 
>>> images = np.array([[[1, 1, 1], [0, 0, 0], [0, 0, 0]] for i in range(4)])
>>> filter = np.array([[1, 1], [-1, -1]])
>>> result = signal.convolve(images, filter, mode='valid', axes=([1, 2], [0, 1]))
>>> result.shape
(4, 2, 2)
>>> result
array([[[-2, -2],
        [ 0,  0]],

       [[-2, -2],
        [ 0,  0]],

       [[-2, -2],
        [ 0,  0]],

       [[-2, -2],
        [ 0,  0]]])

Parameter dot_axes specifies axes on which tensor dot product is being performed. It's also a tuple consisting of a list of dimensions of the first matrix and a list of dimensions of the second matrix. If we specify axes=([], []) it will work exactly as numpy.tensordot, e.g.:

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> c = signal.convolve(a,b, axes=([], []), dot_axes=([1,0], [0,1]))
>>> c.shape
(5, 2)
>>> c
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

By default there is a convolution on all of the axes and tensor dot product on none of them.

Combining those operations can be used e.g. in convolutional neural networks:

import autograd.numpy.random as npr

data = npr.randn(256, 3, 28, 28)
kernel = npr.randn(3, 5, 5)
autograd.scipy.signal.convolve(data, kernel, axes=([2, 3], [1, 2]), dot_axes=([1], [0]))

More examples can be found in tests: https://github.com/BB-UCL/autograd-forward/blob/master/tests/test_scipy.py#L75-L119

[yoavram]

This would be a great addition to the docs, in my opinion.