code: def abs_square(x):
      """Returns the squared value of `tf.abs(x)` for real and complex dtypes."""
      if dtype_util.is_complex(x.dtype):
          return tf.math.square(tf.math.real(x)) + tf.math.square(
              tf.math.imag(x)
          )
      else:
          return tf.math.square(x)
  
file_path: probability/tensorflow_probability/python/math/ode/runge_kutta_util.py
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 2 for calling a global function with example usage: abs_square(x) and returns: unknown, tf