Support for np.gradient

[clemisch]

Is it possible to add support for numpy.gradient in autograd.grad?

I'm using numpy.gradient for computing a cost function in statistical tomographic reconstruction. As it is an optimization task, I need the gradient for the cost function, which is particularly nasty to derive manually when taking into account the pairwise neighbourhoods of pixels. Therefore it would be awesome if I could use autograd for this.

Minimal example:

import autograd.numpy as np
import autograd as ag

def grad_test(x):
    return np.sum(np.gradient(x))

vg = ag.grad(grad_test)
vg(np.random.rand(10, 10, 10))

gives NotImplementedError: VJP of gradient wrt argnums (0,) not defined :cry:

PS: What does the error mean?

[clemisch]

For anyone interested, I found a sort-of workaround for this. You essentially can compute the gradient via autograd.scipy.signal.convolve with a 0.5, 0, -0.5 kernel.

Minimal example:

import autograd as ag
import autograd.scipy.signal as sig
import autograd.numpy as np

# example array to compute the gradient over
arr = np.arange(1000.).reshape(10, 10, 10)

def grad_ag(x):
    """ Compute spacial gradient via convolve """

    # convolve kernel
    grad_kern = np.array([0.5, 0, -0.5])

    return np.array([
        sig.convolve(x, grad_kern[:, None, None], mode='full')[1:-1],
        sig.convolve(x, grad_kern[None, :, None], mode='full')[:, 1:-1],
        sig.convolve(x, grad_kern[None, None, :], mode='full')[..., 1:-1],
    ])

# reference gradient from numpy (can't be autograd'ed)
g_np = np.array(np.gradient(arr))

# manual gradient via convolve
g_ag = grad_ag(arr)

# border handling is different, therefore remove border :(
print np.allclose(g_np[1:-1, 1:-1, 1:-1], g_ag[1:-1, 1:-1, 1:-1])
# => True

def value(x):
    """ Example scalar function using gradient """
    return np.sum(np.square(grad_ag(x)))

# derivative of value
vg = ag.grad(value)

# compute derivative at arr
g = vg(arr)

# It works! :)

I'm still keeping this issue open as having np.gradient would be neater.

PS: Using scipy.ndimage.convolve with mode="nearest" would be better for the border handling (and would perfectly match np.gradient) but sadly it's not implemented in autograd.

[clemisch]

Solved in commit e3b5253