code: def weighted_sum(weights, list_of_states):
      """Computes a weighted sum of `list_of_states`.
  
      Args:
        weights: List of scalar tensors.
        list_of_states: List of states. Every element is assumed to be of the same
          structure of Tensors. Must be of the same length as `weights`.
  
      Returns:
        weighted_sum: A weighted sum of states in `list_of_states`. Has the same
          structure as elements of `list_of_states`.
  
      Raises:
        ValueError: If `list_of_states` is empty or length doesn't match `weights`.
      """
      with tf.name_scope("weighted_sum"):
          if not weights:
              raise ValueError(
                  "`list_of_states` and `weights` must be non-empty"
              )
          if len(weights) != len(list_of_states):
              raise ValueError(
                  "`weights` and `list_of_states` must have same length"
              )
          for state in list_of_states:
              tf.nest.assert_same_structure(state, list_of_states[-1])
          weights_and_states = zip(weights, list_of_states)
          weighted_states = [
              [
                  tf.cast(w, s_comp.dtype) * s_comp
                  for s_comp in tf.nest.flatten(s)
              ]
              for w, s in weights_and_states
              if _possibly_nonzero(w)
          ]
          list_of_components = zip(
              *weighted_states
          )  # Put same components together.
          flat_final_state = [
              tf.add_n(component) for component in list_of_components
          ]
          if not flat_final_state:
              return nest_constant(list_of_states[0], 0.0)
          return tf.nest.pack_sequence_as(
              list_of_states[0], flat_final_state
          )
  
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 3 for calling a global function with example usage: weighted_sum(weights, list_of_states) and returns: tf, nest_constant