2nd order derivs wants insane amount of ram

import jax
import jax.numpy as np
import numpy as onp

def E_fn(conf):
    ri = np.expand_dims(conf, 0)
    rj = np.expand_dims(conf, 1)
    dxdydz = np.power(ri - rj, 2)
    dij = np.sqrt(np.sum(dxdydz, axis=-1))
    return np.sum(dij)

dE_dx_fn = jax.jacrev(E_fn, argnums=(0,))
d2E_dx2_fn = jax.jacfwd(dE_dx_fn, argnums=(0,))

d2E_dx2_fn(onp.random.rand(2483, 3))

Results in:

RuntimeError: Resource exhausted: Out of memory while trying to allocate 551102853132 bytes.

This happens on both CPU and GPU.

There’s no reason this calculation should require 551GB’s worth of ram. The explicit hessian is “only” (24833)^24 bytes=221 MB

[mattjj]

Actually under a jit we never run things op-by-op (though a version of Jax from 2017 did things that way); the values propagated through the python code are abstract ones, basically representing a set of possible arrays, and those aren’t backed by any float values. The abstract values just store shape and dtype, and so they don’t take much memory or cause any FLOPs.

The more I think about it, the more I think XLA should be doing this rewrite optimization for us (and by extension jit should be doing it for you). I’ll raise it with the XLA folks and see what they think. In the worst case, if for some unforeseen reason XLA can’t do this optimization, this is a place where custom ops could help.

[proteneer]

This is an extremely simplified (but representative) repro of a much more complicated set of non-truncated potentials that can’t always be reduced using the law of cosines (but it’s a nice trick to compute the Gramian)