code: def nest_rms_norm(nest):
      """Computes root mean squared norm of nested structure of `Tensor`s.
  
      Args:
        nest: Possibly nested structure of `Tensor`s of which RMS norm is computed.
      Returns:
        norm: Scalar floating tensor equal to the RMS norm of `nest.
      """
      sizes = tf.nest.map_structure(tf.size, nest)
      num_elements = tf.add_n(tf.nest.flatten(sizes))
  
      def averaged_sum_squares(input_tensor):
          num_elements_cast = tf.cast(
              num_elements,
              dtype=dtype_util.real_dtype(input_tensor.dtype),
          )
          return (
              tf.reduce_sum(abs_square(input_tensor))
              / num_elements_cast
          )
  
      squared_sums = tf.nest.map_structure(averaged_sum_squares, nest)
      norm = tf.math.sqrt(tf.add_n(tf.nest.flatten(squared_sums)))
      return norm
  
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 5 for calling a global function with example usage: nest_rms_norm(nest) and returns: norm, unknown