code: def filter_elements(time_indep, time_dep, observation):
      """Convert data into form suitable for filtering with `scan_associative`."""
      with tf.name_scope("filter_elements"):
          f1_unmasked = init_element(
              time_dep.observation_matrix[0],
              time_dep.observation_cov[0],
              time_indep.initial_mean,
              time_indep.initial_cov,
              time_dep.observation_mean[0],
              observation.y[0],
          )
          fk_unmasked = mid_elements(
              time_dep.transition_matrix[:-1],
              time_dep.transition_cov[:-1],
              time_dep.observation_matrix[1:],
              time_dep.observation_cov[1:],
              time_dep.transition_mean[:-1],
              time_dep.observation_mean[1:],
              observation.y[1:],
          )
          elements = tf.nest.map_structure(
              lambda m, ns: tf.concat([[m], ns], axis=0),
              f1_unmasked,
              fk_unmasked,
          )
          if observation.mask is not None:
              f1_masked = init_element_masked(
                  time_indep.initial_mean, time_indep.initial_cov
              )
              fk_masked = mid_elements_masked(
                  time_dep.transition_matrix[:-1],
                  time_dep.transition_cov[:-1],
                  time_dep.transition_mean[:-1],
              )
              elements_masked = tf.nest.map_structure(
                  lambda m, ns: tf.concat([[m], ns], axis=0),
                  f1_masked,
                  fk_masked,
              )
              masks = FilterElements(
                  posterior_link_matrix=observation.mask[
                      ..., None, None
                  ],
                  posterior_mean=observation.mask[..., None],
                  posterior_cov=observation.mask[..., None, None],
                  marginal_likelihood_meanprec=observation.mask[
                      ..., None
                  ],
                  marginal_likelihood_prec=observation.mask[
                      ..., None, None
                  ],
              )
              elements = tf.nest.map_structure(
                  tf.where, masks, elements_masked, elements
              )
  
      return elements
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math import linalg
  from tensorflow_probability.python.math.scan_associative import scan_associative
label_desc: Use image node_id 9 for calling a global function with example usage: filter_elements(time_indep, time_dep, observation) and returns: elements