code: def _write_solution_components(solution, solutions_arrays, time_id):
      """Writes individual `Tensor` components of `solution` to `solutions_arrays`.
  
      Args:
        solution: Possibly nested structure of `Tensor`s representing current
          solution to be written to `solutions_arrays`.
        solutions_arrays: List of `TensorArray`s where `solution` components are
          written. Must have the same number of elements as components in
          `solution`.
        time_id: Integer `Tensor` representing the index at which `solution`
          components are written to `solutions_arrays`.
  
      Returns:
        updated_arrays: List of `TensorArray`s whose components at index `time_id`
          now contain components of the `solution`.
      """
      tf.nest.assert_same_structure(solution, solutions_arrays)
      write_solution = lambda array, tensor: array.write(
          time_id, tensor
      )
      updated_arrays = tf.nest.map_structure(
          write_solution, solutions_arrays, solution
      )
      return updated_arrays
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.ode import base
  from tensorflow_probability.python.math.ode import runge_kutta_util
  from tensorflow_probability.python.math.ode import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 4 for calling a global function with example usage: _write_solution_components(solution, solutions_arrays, time_id) and returns: updated_arrays