def elementwise_grad(fun, argnum=0):
    def sum_output(*args, **kwargs):
        return np.sum(fun(*args, **kwargs))
    return grad(sum_output, argnum=argnum)
Like `jacobian`, but produces a function which computes just the diagonal
dataset/ETHPy150Open HIPS/autograd/autograd/convenience_wrappers.py/elementwise_grad