def _non_nan(x): # pylint: disable=invalid-name
"""Replaces NaN values with zeros.
A support function replaces NaN values with zeros to escape
the undefined behavior of the division by zero.
Args:
x: tensor with arbitrary shape.
Returns:
Array with NaNs replaced with 0.
"""
return jnp.where(jnp.isnan(x), 0., x)
def non_nan(x): # pylint: disable=invalid-name
return jnp.where(jnp.isnan(x), 0., x)