def build_kalman_filter_step(get_transition_matrix_for_timestep,<EOL>get_transition_noise_for_timestep,<EOL>get_observation_matrix_for_timestep,<EOL>get_observation_noise_for_timestep):
def kalman_filter_step(state, elems_t):<EOL><INDENT>"""<STR_LIT>"""<EOL>if isinstance(elems_t, tuple):<EOL><INDENT>x_t, mask_t = elems_t<EOL><DEDENT>else:<EOL><INDENT>x_t = elems_t<EOL>mask_t = None<EOL><DEDENT>observation_matrix = get_observation_matrix_for_timestep(state.timestep)<EOL>observation_noise = get_observation_noise_for_timestep(state.timestep)<EOL>if mask_t is not None:<EOL><INDENT>x_expected = _propagate_mean(state.predicted_mean,<EOL>observation_matrix,<EOL>observation_noise) * tf.ones_like(x_t)<EOL>x_t = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=x_expected)), x_expected,<EOL>tf.broadcast_to(x_t, tf.shape(input=x_expected)))<EOL><DEDENT>(filtered_mean,<EOL>filtered_cov,<EOL>observation_dist) = linear_gaussian_update(<EOL>state.predicted_mean, state.predicted_cov,<EOL>observation_matrix, observation_noise,<EOL>x_t)<EOL>log_marginal_likelihood = observation_dist.log_prob(x_t[..., <NUM_LIT:0>])<EOL>if mask_t is not None:<EOL><INDENT>filtered_mean = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=filtered_mean)),<EOL>state.predicted_mean, filtered_mean)<EOL>filtered_cov = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=filtered_cov)),<EOL>state.predicted_cov, filtered_cov)<EOL>log_marginal_likelihood = tf.where(<EOL>tf.broadcast_to(mask_t[..., <NUM_LIT:0>, <NUM_LIT:0>],<EOL>tf.shape(input=log_marginal_likelihood)),<EOL>tf.zeros_like(log_marginal_likelihood),<EOL>log_marginal_likelihood)<EOL><DEDENT>predicted_mean, predicted_cov = kalman_transition(<EOL>filtered_mean,<EOL>filtered_cov,<EOL>get_transition_matrix_for_timestep(state.timestep),<EOL>get_transition_noise_for_timestep(state.timestep))<EOL>return KalmanFilterState(<EOL>filtered_mean, filtered_cov,<EOL>predicted_mean, predicted_cov,<EOL>observation_dist.mean()[..., tf.newaxis],<EOL>observation_dist.covariance(),<EOL>log_marginal_likelihood,<EOL>state.timestep+<NUM_LIT:1>)<EOL><DEDENT>return kalman_filter_step<EOL>
Build a callable that performs one step of Kalman filtering.

    Args:
      get_transition_matrix_for_timestep: callable taking a timestep
        as an integer `Tensor` argument, and returning a `LinearOperator`
        of shape `[latent_size, latent_size]`.
      get_transition_noise_for_timestep: callable taking a timestep as
        an integer `Tensor` argument, and returning a
        `MultivariateNormalLinearOperator` of event shape
        `[latent_size]`.
      get_observation_matrix_for_timestep: callable taking a timestep
        as an integer `Tensor` argument, and returning a `LinearOperator`
        of shape `[observation_size, observation_size]`.
      get_observation_noise_for_timestep: callable taking a timestep as
        an integer `Tensor` argument, and returning a
        `MultivariateNormalLinearOperator` of event shape
        `[observation_size]`.

    Returns:
      kalman_filter_step: a callable that updates a KalmanFilterState
        from timestep `t-1` to `t`.