code: def _write_accumulated_quantities(
      step, accumulated_quantities, updated_estimate
  ):
      """Update arrays passed through the filter loop."""
      new_accumulated_quantities = [
          nest.map_structure_up_to(
              element,
              lambda x, y: x.write(step, y[i]),
              element,
              updated_estimate,
          )  # pylint: disable=cell-var-from-loop
          for i, element in enumerate(accumulated_quantities)
      ]
      return linear_gaussian_ssm.KalmanFilterState(
          *new_accumulated_quantities
      )
  
file_path: probability/tensorflow_probability/python/experimental/sequential/extended_kalman_filter.py
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import linalg
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: _write_accumulated_quantities(step, accumulated_quantities, updated_estimate) and returns: linear_gaussian_ssm