code: 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.0, x)
  
imports: from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
label_desc: Use image node_id 3 for calling a global function with example usage: _non_nan(x) and returns: jnp