Calculating hessian block by block

[cagrikymk]

I am trying to calculate a hessian of a fairly complex function and a simple jax.hessian call is causing out-of-memory error.
My solution is to calculate the hessian block by block by utilizing indexing in a loop which works okay but little slow (slower than the numerical approximation).

Row by row hessian calculation (using a loop): 55 sec
Using finite difference method (from grad -> hessian): 37 sec

Here is a sample code demonstrating the logic (I am using vmap to make a fairer comparison but in the actual implementation I use a for loop to not get OOM):

import jax
import jax.numpy as jnp
from jax import random

def func(x):
    return jnp.sum(x + x[:, None]**2)

def full_hessian(x):
    return jax.hessian(func)(x)

def partial_hessian(i, x):
    def inner(*args):
        return jax.jacrev(func)(*args)[i]
    return jax.jacfwd(inner)(x)


N = 1000
key = random.PRNGKey(0)
x = random.uniform(key, shape=(N,))

j_full_hessian = jax.jit(full_hessian)
j_partial_hessian = jax.jit(jax.vmap(partial_hessian, (0, None)))


h1 = j_full_hessian(x).block_until_ready()
h2 = j_partial_hessian(jnp.arange(N), x).block_until_ready()

print("full_hessian == partial_hessian:", jnp.all(h1 == h2))


%timeit j_full_hessian(x).block_until_ready()
#60.9 µs ± 1.8 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit j_partial_hessian1(jnp.arange(N), x).block_until_ready()
#164 µs ± 2.78 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)