'Understanding autograd's elementwise_grad

I have trouble understanding elementwise_grad. It is supposed to give the sum of each column of the Jacobian ...

from autograd import elementwise_grad,jacobian
import autograd.numpy as np

def testfun1(a1):
    return np.asarray(list(range(1,11))).T * a1

elementwise_grad(testfun1,0)(2*np.ones(shape=(10)))

gives

array([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.])

which I fully understand. The 10x10 Jacobian is diagonal and the sum of the columns is the diagonal.

What I do not get is this:

def testfun2(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
    return np.asarray(list(range(1,11))).T * np.asarray([a1, a2, a3, a4, a5, a6, a7, a8, a9, a10])

testfun2(*(2*np.ones(shape=(10))))

yields

array(0.)

I thought this should have a 10x1 Jacobian (summed over the single column for the output) and yield

array(0.)

I would have expected that

np.asarray([elementwise_grad(testfun2,i)(*(1*np.ones(shape=(10)))) for i in range(10)])

has the same result as elementwise_grad on testfun1 ...



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source