Elementwise Hessian

I’m super impressed with the progress you’ve been making! Autograd has the potential to save me and many others a lot of grief in terms of automatically differentiating scalar-valued functions while retaining broadcasting support. So right now I can do something like

from autograd import elementwise_grad
import numpy as np
def wrapper(args):
    return energy(*args)
grad = elementwise_grad(wrapper)
inp = [1000*np.random.rand(10,1), np.random.rand(1,10), np.random.rand(1,10), np.random.rand(1,10), np.random.rand(1,10), np.random.rand(1,10)]
inp = np.broadcast_arrays(*inp)
%time gres = grad(inp)
print([i.shape for i in gres])

CPU times: user 301 ms, sys: 18 ms, total: 319 ms
Wall time: 300 ms
[(10, 10), (10, 10), (10, 10), (10, 10), (10, 10), (10, 10)]

And so each element of the returned list is the gradient of the objective function, with respect to the variable at that argument index, for all 100 combinations of input variables. So what I want now is the Hessian of my objective function, elementwise, meaning if my gradient is a (6, 10, 10) array, I want back a (6, 6, 10, 10) array with all the second derivatives including cross terms. Is this possible to do with Autograd?

[mattjj]

I think it works. Here’s a cleaner elementwise_hess function for you, as well as a better test. (Of course, don’t trust me, please test it yourself!)

from autograd import elementwise_grad
import autograd.numpy as np

def energy(args):
    return 8.314*args[0]*(args[1]*np.log(args[1]) + args[2]*np.log(args[2]))

gradfun = elementwise_grad(energy)
N = 100
M = 50

inp = [1000*np.random.rand(N,1), np.random.rand(1,M), np.random.rand(1,M)]
inp = np.broadcast_arrays(*inp)


gres = gradfun(inp)
print([i.shape for i in gres])

### here's the new stuff!!

from autograd import jacobian

def elementwise_hess(fun, argnum=0):
    sum_latter_dims = lambda x: np.sum(x.reshape(x.shape[0], -1), 1)
    def sum_grad_output(*args, **kwargs):
        return sum_latter_dims(elementwise_grad(fun)(*args, **kwargs))
    return jacobian(sum_grad_output, argnum)

# make inp an array so that jacobian() can handle it
array_inp = np.concatenate([arr[None,...] for arr in inp])

# check shapes on the given energy function
hessfun = elementwise_hess(energy)
print(hessfun(array_inp).shape)

# test on some scalar elements
from autograd import grad
hessval = hessfun(array_inp)
print grad(lambda x: grad(energy)(x)[1])(array_inp[:,8,42])
print hessval[:,1,8,42]

It prints something like this (different every time due to the pseudorandom numbers generated):

[(100, 50), (100, 50), (100, 50)]
(3, 3, 100, 50)
[   7.10510658  920.2143044     0.        ]
[   7.10510658  920.2143044     0.        ]

Please reopen the issue if it’s not right!

[richardotis]

Actually the above turned out not to be entirely true. There’s still a broadcasting hangup but it’s pretty easily solved. I’m working on a code writeup which I’ll share after I’ve kicked the tires a bit.