NaNs with AD

[OlaRonning]

Hi,

I'm working on an MD project, and I came across an oddity with AD in Jax. The gradient/Jacobian of a function with a jnp.where and jnp.sqrt produces NaNs depending on whether the jnp.sqrt is nested inside the jnp.where or after the jnp.where. If I take the element-wise product with another function and compute the gradient, the opposite pair of functions produces NaNs. I have a minimal example below.

from pprint import pprint

import jax.numpy as jnp
from jax import jacfwd, device_get, grad


def nested_sqrt(cos_angle1):
    energy = jnp.where(cos_angle1 > 0., jnp.sqrt(cos_angle1), 0.)  # sqrt here
    return energy


def outer_sqrt(cos_angle1):
    energy = jnp.where(cos_angle1 > 0., cos_angle1, 0.)
    return jnp.sqrt(energy)  # sqrt here


def simp_potent(dist):
    return dist


if __name__ == '__main__':
    cos_angle = jnp.array([0.01866167, -0.15954296])
    dists = jnp.array([2.5986805, 6.3012958])

    print('Kinetic (nested sqrt)')  # No NaNs
    pprint(device_get(jacfwd(nested_sqrt)(cos_angle).tolist()))

    print('Total energy (nested sqrt):')  # Produces NaNs
    pprint(device_get(grad(lambda a, d: (nested_sqrt(a) * simp_potent(d)).sum())(cos_angle, dists)).tolist())

    print('Kinetic (outer sqrt)')  # Produces NaNs
    pprint(device_get(jacfwd(outer_sqrt)(cos_angle).tolist()))

    print('Total energy (outer sqrt):')  # No NaNs
    pprint(device_get(grad(lambda a, d: (outer_sqrt(a) * simp_potent(d)).sum())(cos_angle, dists)).tolist())

I assume the difference in their behavior arises from the way the IR is generated; however, I've not found the exact cause. What is the cause of this? And, is the behavior expected?

[jakevdp]

Thanks for the question! This is a commonly-encountered issue in autodiff systems; you can read about it here: FAQ: Gradients contain NaN where using where

[OlaRonning]

Thanks, that clarifies it.