def _augment_sample_shape(partial_batch_dist,
                          full_sample_and_batch_shape,
                          validate_args=False):
  """Augment a sample shape to broadcast batch dimensions.

  Computes an augmented sample shape, so that any batch dimensions not
  part of the distribution `partial_batch_dist` are treated as identical
  distributions.

  # partial_batch_dist.batch_shape  = [      7]
  # full_sample_and_batch_shape     = [3, 4, 7]
  # => return an augmented sample shape of [3, 4] so that
  #    partial_batch_dist.sample(augmented_sample_shape) has combined
  #    sample and batch shape of [3, 4, 7].

  Args:
    partial_batch_dist: `tfd.Distribution` instance with batch shape a
      prefix of `full_sample_and_batch_shape`.
    full_sample_and_batch_shape: a Tensor or Tensor-like shape.
    validate_args: if True, check for shape errors at runtime.
  Returns:
    augmented_sample_shape: sample shape such that
      `partial_batch_dist.sample(augmented_sample_shape)` has combined
      sample and batch shape of `full_sample_and_batch_shape`.

  Raises:
    ValueError: if `partial_batch_dist.batch_shape` has more dimensions than
      `full_sample_and_batch_shape`.
    NotImplementedError: if broadcasting would be required to make
      `partial_batch_dist.batch_shape` into a prefix of
      `full_sample_and_batch_shape` .
  """
  full_ndims = distribution_util.prefer_static_shape(
      full_sample_and_batch_shape)[0]
  partial_batch_ndims = (
      tensorshape_util.rank(partial_batch_dist.batch_shape)  # pylint: disable=g-long-ternary
      if tensorshape_util.rank(partial_batch_dist.batch_shape) is not None
      else distribution_util.prefer_static_shape(
          partial_batch_dist.batch_shape_tensor())[0])

  num_broadcast_dims = full_ndims - partial_batch_ndims

  expected_partial_batch_shape = (
      full_sample_and_batch_shape[num_broadcast_dims:])
  expected_partial_batch_shape_static = tf.get_static_value(
      full_sample_and_batch_shape[num_broadcast_dims:])

  # Raise errors statically if possible.
  num_broadcast_dims_static = tf.get_static_value(num_broadcast_dims)
  if num_broadcast_dims_static is not None:
    if num_broadcast_dims_static < 0:
      raise ValueError("Cannot broadcast distribution {} batch shape to "
                       "target batch shape with fewer dimensions"
                       .format(partial_batch_dist))
  if (expected_partial_batch_shape_static is not None and
      tensorshape_util.is_fully_defined(partial_batch_dist.batch_shape)):
    if (partial_batch_dist.batch_shape and
        any(expected_partial_batch_shape_static != tensorshape_util.as_list(
            partial_batch_dist.batch_shape))):
      raise NotImplementedError("Broadcasting is not supported; "
                                "unexpected batch shape "
                                "(expected {}, saw {}).".format(
                                    expected_partial_batch_shape_static,
                                    partial_batch_dist.batch_shape
                                ))
  runtime_assertions = []
  if validate_args:
    runtime_assertions.append(
        assert_util.assert_greater_equal(
            tf.convert_to_tensor(value=num_broadcast_dims, dtype=tf.int32),
            tf.zeros((), dtype=tf.int32),
            message=("Cannot broadcast distribution {} batch shape to "
                     "target batch shape with fewer dimensions.".format(
                         partial_batch_dist))))
    runtime_assertions.append(
        assert_util.assert_equal(
            expected_partial_batch_shape,
            partial_batch_dist.batch_shape_tensor(),
            message=("Broadcasting is not supported; "
                     "unexpected batch shape."),
            name="assert_batch_shape_same"))

  with tf.control_dependencies(runtime_assertions):
    return full_sample_and_batch_shape[:num_broadcast_dims]

Augment a sample shape to broadcast batch dimensions .


def build_backward_pass_step(get_transition_matrix_for_timestep):
  """Build a callable that perform one step for backward smoothing.

  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]`.

  Returns:
    backward_pass_step: a callable that updates a BackwardPassState
      from timestep `t` to `t-1`.
  """

  def backward_pass_step(state,
                         filtered_parameters):
    """Run a single step of backward smoothing."""

    (filtered_mean, filtered_cov,
     predicted_mean, predicted_cov) = filtered_parameters
    transition_matrix = get_transition_matrix_for_timestep(state.timestep)

    next_posterior_mean = state.backward_mean
    next_posterior_cov = state.backward_cov

    posterior_mean, posterior_cov = backward_smoothing_update(
        filtered_mean,
        filtered_cov,
        predicted_mean,
        predicted_cov,
        next_posterior_mean,
        next_posterior_cov,
        transition_matrix)

    return BackwardPassState(backward_mean=posterior_mean,
                             backward_cov=posterior_cov,
                             timestep=state.timestep-1)

  return backward_pass_step

Build a callable that perform one step for backward smoothing .


def backward_smoothing_update(filtered_mean,
                              filtered_cov,
                              predicted_mean,
                              predicted_cov,
                              next_posterior_mean,
                              next_posterior_cov,
                              transition_matrix):
  """Backward update for a Kalman smoother.

  Give the `filtered_mean` mu(t | t), `filtered_cov` sigma(t | t),
  `predicted_mean` mu(t+1 | t) and `predicted_cov` sigma(t+1 | t),
  as returns from the `forward_filter` function, as well as
  `next_posterior_mean` mu(t+1 | 1:T) and `next_posterior_cov` sigma(t+1 | 1:T),
  if the `transition_matrix` of states from time t to time t+1
  is given as A(t+1), the 1 step backward smoothed distribution parameter
  could be calculated as:
  p(z(t) | Obs(1:T)) = N( mu(t | 1:T), sigma(t | 1:T)),
  mu(t | 1:T) = mu(t | t) + J(t) * (mu(t+1 | 1:T) - mu(t+1 | t)),
  sigma(t | 1:T) = sigma(t | t)
                   + J(t) * (sigma(t+1 | 1:T) - sigma(t+1 | t) * J(t)',
  J(t) = sigma(t | t) * A(t+1)' / sigma(t+1 | t),
  where all the multiplications are matrix multiplication, and `/` is
  the matrix inverse. J(t) is the backward Kalman gain matrix.

  The algorithm can be intialized from mu(T | 1:T) and sigma(T | 1:T),
  which are the last step parameters returned by forward_filter.

  Args:
    filtered_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`, containing mu(t | t).
    filtered_cov: `Tensor` with event shape `[latent_size, latent_size]` and
      batch shape `B`, containing sigma(t | t).
    predicted_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`, containing mu(t+1 | t).
    predicted_cov: `Tensor` with event shape `[latent_size, latent_size]` and
      batch shape `B`, containing sigma(t+1 | t).
    next_posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`, containing mu(t+1 | 1:T).
    next_posterior_cov: `Tensor` with event shape `[latent_size, latent_size]`
      and batch shape `B`, containing sigma(t+1 | 1:T).
    transition_matrix: `LinearOperator` with shape
      `[latent_size, latent_size]` and batch shape broadcastable
      to `B`.

  Returns:
    posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`, containing mu(t | 1:T).
    posterior_cov: `Tensor` with event shape `[latent_size, latent_size]` and
      batch shape `B`, containing sigma(t | 1:T).
  """
  # Compute backward Kalman gain:
  # J = F * T' * P^{-1}
  # Since both F(iltered) and P(redictive) are cov matrices,
  # thus self-adjoint, we can take the transpose.
  # computation:
  #      = (P^{-1} * T * F)'
  #      = (P^{-1} * tmp_gain_cov) '
  #      = (P \ tmp_gain_cov)'
  tmp_gain_cov = transition_matrix.matmul(filtered_cov)
  predicted_cov_chol = tf.linalg.cholesky(predicted_cov)
  gain_transpose = tf.linalg.cholesky_solve(predicted_cov_chol, tmp_gain_cov)

  posterior_mean = (filtered_mean +
                    tf.linalg.matmul(gain_transpose,
                                     next_posterior_mean - predicted_mean,
                                     adjoint_a=True))
  posterior_cov = (
      filtered_cov +
      tf.linalg.matmul(gain_transpose,
                       tf.linalg.matmul(
                           next_posterior_cov - predicted_cov, gain_transpose),
                       adjoint_a=True))

  return (posterior_mean, posterior_cov)

def backward_smoothing_update ( filtered_mean , filtered_cov , predicted_mean , predicted_cov , next_posterior_mean , next_posterior_cov , transition_matrix ) : tmp_gain_cov = transition_matrix . matmul ( filtered_cov ) predicted_cov_chol = tf . linalg . cholesky ( predicted_cov ) gain_transpose = tf . linalg . cholesky_solve ( predicted_cov_chol , tmp_gain_cov ) posterior_mean = ( filtered_mean + tf . linalg . matmul ( gain_transpose , next_posterior_mean - predicted_mean , adjoint_a = True ) ) posterior_cov = ( filtered_cov + tf . linalg . matmul ( gain_transpose , tf . linalg . matmul ( next_posterior_cov - predicted_cov , gain_transpose ) , adjoint_a = True ) ) return ( posterior_mean , posterior_cov )


def build_kalman_filter_step(get_transition_matrix_for_timestep,
                             get_transition_noise_for_timestep,
                             get_observation_matrix_for_timestep,
                             get_observation_noise_for_timestep):
  """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`.
  """

  def kalman_filter_step(state, elems_t):
    """Run a single step of Kalman filtering.

    Args:
      state: A `KalmanFilterState` object representing the previous
        filter state at time `t-1`.
      elems_t: A tuple of Tensors `(x_t, mask_t)`, or a `Tensor` `x_t`.
        `x_t` is a `Tensor` with rightmost shape dimensions
        `[observation_size, 1]` representing the vector observed at time `t`,
        and `mask_t` is a `Tensor` with rightmost dimensions`[1, 1]`
        representing the observation mask at time `t`. Both `x_t` and `mask_t`
        may have batch dimensions, which must be compatible with the batch
        dimensions of `state.predicted_mean` and `state.predictived_cov`
        respectively. If `mask_t` is not provided, it is assumed to be `None`.

    Returns:
      new_state: A `KalmanFilterState` object representing the new
        filter state at time `t`.
    """

    if isinstance(elems_t, tuple):
      x_t, mask_t = elems_t
    else:
      x_t = elems_t
      mask_t = None

    observation_matrix = get_observation_matrix_for_timestep(state.timestep)
    observation_noise = get_observation_noise_for_timestep(state.timestep)
    if mask_t is not None:
      # Before running the update, fill in masked observations using the prior
      # expectation. The precise filled value shouldn't matter since updates
      # from masked elements will not be selected below, but we need to ensure
      # that any results we incidently compute on masked values are at least
      # finite (not inf or NaN) so that they don't screw up gradient propagation
      # through `tf.where`, as described in
      #  https://github.com/tensorflow/tensorflow/issues/2540.
      # We fill with the prior expectation because any fixed value such as zero
      # might be arbitrarily unlikely under the prior, leading to overflow in
      # the updates, but the prior expectation should always be a
      # 'reasonable' observation.
      x_expected = _propagate_mean(state.predicted_mean,
                                   observation_matrix,
                                   observation_noise) * tf.ones_like(x_t)
      x_t = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=x_expected)), x_expected,
          tf.broadcast_to(x_t, tf.shape(input=x_expected)))

    # Given predicted mean u_{t|t-1} and covariance P_{t|t-1} from the
    # previous step, incorporate the observation x_t, producing the
    # filtered mean u_t and covariance P_t.
    (filtered_mean,
     filtered_cov,
     observation_dist) = linear_gaussian_update(
         state.predicted_mean, state.predicted_cov,
         observation_matrix, observation_noise,
         x_t)

    # Compute the marginal likelihood p(x_{t} | x_{:t-1}) for this
    # observation.
    log_marginal_likelihood = observation_dist.log_prob(x_t[..., 0])

    if mask_t is not None:
      filtered_mean = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_mean)),
          state.predicted_mean, filtered_mean)
      filtered_cov = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_cov)),
          state.predicted_cov, filtered_cov)
      log_marginal_likelihood = tf.where(
          tf.broadcast_to(mask_t[..., 0, 0],
                          tf.shape(input=log_marginal_likelihood)),
          tf.zeros_like(log_marginal_likelihood),
          log_marginal_likelihood)

    # Run the filtered posterior through the transition
    # model to predict the next time step:
    #  u_{t|t-1} = F_t u_{t-1} + b_t
    #  P_{t|t-1} = F_t P_{t-1} F_t' + Q_t
    predicted_mean, predicted_cov = kalman_transition(
        filtered_mean,
        filtered_cov,
        get_transition_matrix_for_timestep(state.timestep),
        get_transition_noise_for_timestep(state.timestep))

    return KalmanFilterState(
        filtered_mean, filtered_cov,
        predicted_mean, predicted_cov,
        observation_dist.mean()[..., tf.newaxis],
        observation_dist.covariance(),
        log_marginal_likelihood,
        state.timestep+1)

  return kalman_filter_step

def build_kalman_filter_step ( get_transition_matrix_for_timestep , get_transition_noise_for_timestep , get_observation_matrix_for_timestep , get_observation_noise_for_timestep ) : def kalman_filter_step ( state , elems_t ) : if isinstance ( elems_t , tuple ) : x_t , mask_t = elems_t else : x_t = elems_t mask_t = None observation_matrix = get_observation_matrix_for_timestep ( state . timestep ) observation_noise = get_observation_noise_for_timestep ( state . timestep ) if mask_t is not None : x_expected = _propagate_mean ( state . predicted_mean , observation_matrix , observation_noise ) * tf . ones_like ( x_t ) x_t = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = x_expected ) ) , x_expected , tf . broadcast_to ( x_t , tf . shape ( input = x_expected ) ) ) ( filtered_mean , filtered_cov , observation_dist ) = linear_gaussian_update ( state . predicted_mean , state . predicted_cov , observation_matrix , observation_noise , x_t ) log_marginal_likelihood = observation_dist . log_prob ( x_t [ ... , 0 ] ) if mask_t is not None : filtered_mean = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = filtered_mean ) ) , state . predicted_mean , filtered_mean ) filtered_cov = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = filtered_cov ) ) , state . predicted_cov , filtered_cov ) log_marginal_likelihood = tf . where ( tf . broadcast_to ( mask_t [ ... , 0 , 0 ] , tf . shape ( input = log_marginal_likelihood ) ) , tf . zeros_like ( log_marginal_likelihood ) , log_marginal_likelihood ) predicted_mean , predicted_cov = kalman_transition ( filtered_mean , filtered_cov , get_transition_matrix_for_timestep ( state . timestep ) , get_transition_noise_for_timestep ( state . timestep ) ) return KalmanFilterState ( filtered_mean , filtered_cov , predicted_mean , predicted_cov , observation_dist . mean ( ) [ ... , tf . newaxis ] , observation_dist . covariance ( ) , log_marginal_likelihood , state . timestep + 1 ) return kalman_filter_step

Build a callable that performs one step of Kalman filtering .


def linear_gaussian_update(
    prior_mean, prior_cov, observation_matrix, observation_noise, x_observed):
  """Conjugate update for a linear Gaussian model.

  Given a normal prior on a latent variable `z`,
    `p(z) = N(prior_mean, prior_cov) = N(u, P)`,
  for which we observe a linear Gaussian transformation `x`,
    `p(x|z) = N(H * z + c, R)`,
  the posterior is also normal:
    `p(z|x) = N(u*, P*)`.

  We can write this update as
     x_expected = H * u + c # pushforward prior mean
     S = R + H * P * H'  # pushforward prior cov
     K = P * H' * S^{-1} # optimal Kalman gain
     u* = u + K * (x_observed - x_expected) # posterior mean
     P* = (I - K * H) * P (I - K * H)' + K * R * K' # posterior cov
  (see, e.g., https://en.wikipedia.org/wiki/Kalman_filter#Update)

  Args:
    prior_mean: `Tensor` with event shape `[latent_size, 1]` and
      potential batch shape `B = [b1, ..., b_n]`.
    prior_cov: `Tensor` with event shape `[latent_size, latent_size]`
      and batch shape `B` (matching `prior_mean`).
    observation_matrix: `LinearOperator` with shape
      `[observation_size, latent_size]` and batch shape broadcastable
      to `B`.
    observation_noise: potentially-batched
      `MultivariateNormalLinearOperator` instance with event shape
      `[observation_size]` and batch shape broadcastable to `B`.
    x_observed: potentially batched `Tensor` with event shape
      `[observation_size, 1]` and batch shape `B`.

  Returns:
    posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`.
    posterior_cov: `Tensor` with event shape `[latent_size,
      latent_size]` and batch shape `B`.
    predictive_dist: the prior predictive distribution `p(x|z)`,
      as a `Distribution` instance with event
      shape `[observation_size]` and batch shape `B`. This will
      typically be `tfd.MultivariateNormalTriL`, but when
      `observation_size=1` we return a `tfd.Independent(tfd.Normal)`
      instance as an optimization.
  """

  # If observations are scalar, we can avoid some matrix ops.
  observation_size_is_static_and_scalar = (
      tf.compat.dimension_value(observation_matrix.shape[-2]) == 1)

  # Push the predicted mean for the latent state through the
  # observation model
  x_expected = _propagate_mean(prior_mean,
                               observation_matrix,
                               observation_noise)

  # Push the predictive covariance of the latent state through the
  # observation model:
  #  S = R + H * P * H'.
  # We use a temporary variable for H * P,
  # reused below to compute Kalman gain.
  tmp_obs_cov = observation_matrix.matmul(prior_cov)
  predicted_obs_cov = (
      observation_matrix.matmul(tmp_obs_cov, adjoint_arg=True)
      + observation_noise.covariance())

  # Compute optimal Kalman gain:
  #  K = P * H' * S^{-1}
  # Since both S and P are cov matrices, thus symmetric,
  # we can take the transpose and reuse our previous
  # computation:
  #      = (S^{-1} * H * P)'
  #      = (S^{-1} * tmp_obs_cov) '
  #      = (S \ tmp_obs_cov)'
  if observation_size_is_static_and_scalar:
    gain_transpose = tmp_obs_cov/predicted_obs_cov
  else:
    predicted_obs_cov_chol = tf.linalg.cholesky(predicted_obs_cov)
    gain_transpose = tf.linalg.cholesky_solve(predicted_obs_cov_chol,
                                              tmp_obs_cov)

  # Compute the posterior mean, incorporating the observation.
  #  u* = u + K (x_observed - x_expected)
  posterior_mean = (prior_mean +
                    tf.linalg.matmul(gain_transpose, x_observed - x_expected,
                                     adjoint_a=True))

  # For the posterior covariance, we could use the simple update
  #  P* = P - K * H * P
  # but this is prone to numerical issues because it subtracts a
  # value from a PSD matrix.  We choose instead to use the more
  # expensive Jordan form update
  #  P* = (I - K H) * P * (I - K H)' + K R K'
  # which always produces a PSD result. This uses
  #  tmp_term = (I - K * H)'
  # as an intermediate quantity.
  tmp_term = -observation_matrix.matmul(gain_transpose, adjoint=True)  # -K * H
  tmp_term = tf.linalg.set_diag(tmp_term, tf.linalg.diag_part(tmp_term) + 1)
  posterior_cov = (
      tf.linalg.matmul(
          tmp_term, tf.linalg.matmul(prior_cov, tmp_term), adjoint_a=True)
      + tf.linalg.matmul(gain_transpose,
                         tf.linalg.matmul(
                             observation_noise.covariance(), gain_transpose),
                         adjoint_a=True))

  if observation_size_is_static_and_scalar:
    # A plain Normal would have event shape `[]`; wrapping with Independent
    # ensures `event_shape=[1]` as required.
    predictive_dist = independent.Independent(
        normal.Normal(loc=x_expected[..., 0],
                      scale=tf.sqrt(predicted_obs_cov[..., 0])),
        reinterpreted_batch_ndims=1)

    # Minor hack to define the covariance, so that `predictive_dist` can pass as
    # an MVNTriL-like object.
    predictive_dist.covariance = lambda: predicted_obs_cov
  else:
    predictive_dist = mvn_tril.MultivariateNormalTriL(
        loc=x_expected[..., 0],
        scale_tril=predicted_obs_cov_chol)

  return posterior_mean, posterior_cov, predictive_dist

def linear_gaussian_update ( prior_mean , prior_cov , observation_matrix , observation_noise , x_observed ) : observation_size_is_static_and_scalar = ( tf . compat . dimension_value ( observation_matrix . shape [ - 2 ] ) == 1 ) x_expected = _propagate_mean ( prior_mean , observation_matrix , observation_noise ) tmp_obs_cov = observation_matrix . matmul ( prior_cov ) predicted_obs_cov = ( observation_matrix . matmul ( tmp_obs_cov , adjoint_arg = True ) + observation_noise . covariance ( ) ) if observation_size_is_static_and_scalar : gain_transpose = tmp_obs_cov / predicted_obs_cov else : predicted_obs_cov_chol = tf . linalg . cholesky ( predicted_obs_cov ) gain_transpose = tf . linalg . cholesky_solve ( predicted_obs_cov_chol , tmp_obs_cov ) posterior_mean = ( prior_mean + tf . linalg . matmul ( gain_transpose , x_observed - x_expected , adjoint_a = True ) ) tmp_term = - observation_matrix . matmul ( gain_transpose , adjoint = True ) tmp_term = tf . linalg . set_diag ( tmp_term , tf . linalg . diag_part ( tmp_term ) + 1 ) posterior_cov = ( tf . linalg . matmul ( tmp_term , tf . linalg . matmul ( prior_cov , tmp_term ) , adjoint_a = True ) + tf . linalg . matmul ( gain_transpose , tf . linalg . matmul ( observation_noise . covariance ( ) , gain_transpose ) , adjoint_a = True ) ) if observation_size_is_static_and_scalar : predictive_dist = independent . Independent ( normal . Normal ( loc = x_expected [ ... , 0 ] , scale = tf . sqrt ( predicted_obs_cov [ ... , 0 ] ) ) , reinterpreted_batch_ndims = 1 ) predictive_dist . covariance = lambda : predicted_obs_cov else : predictive_dist = mvn_tril . MultivariateNormalTriL ( loc = x_expected [ ... , 0 ] , scale_tril = predicted_obs_cov_chol ) return posterior_mean , posterior_cov , predictive_dist

Conjugate update for a linear Gaussian model .


def kalman_transition(filtered_mean, filtered_cov,
                      transition_matrix, transition_noise):
  """Propagate a filtered distribution through a transition model."""

  predicted_mean = _propagate_mean(filtered_mean,
                                   transition_matrix,
                                   transition_noise)
  predicted_cov = _propagate_cov(filtered_cov,
                                 transition_matrix,
                                 transition_noise)
  return predicted_mean, predicted_cov

Propagate a filtered distribution through a transition model .


def build_kalman_mean_step(get_transition_matrix_for_timestep,
                           get_transition_noise_for_timestep,
                           get_observation_matrix_for_timestep,
                           get_observation_noise_for_timestep):
  """Build a callable that performs one step of Kalman mean recursion.

  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_mean_step: a callable that computes latent state and
      observation means at time `t`, given latent mean at time `t-1`.
  """

  def mean_step(previous_means, t):
    """Single step of prior mean recursion."""
    previous_latent_mean, _ = previous_means

    latent_mean = _propagate_mean(previous_latent_mean,
                                  get_transition_matrix_for_timestep(t - 1),
                                  get_transition_noise_for_timestep(t - 1))
    observation_mean = _propagate_mean(latent_mean,
                                       get_observation_matrix_for_timestep(t),
                                       get_observation_noise_for_timestep(t))
    return (latent_mean, observation_mean)

  return mean_step

Build a callable that performs one step of Kalman mean recursion .


def build_kalman_cov_step(get_transition_matrix_for_timestep,
                          get_transition_noise_for_timestep,
                          get_observation_matrix_for_timestep,
                          get_observation_noise_for_timestep):
  """Build a callable for one step of Kalman covariance recursion.

  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:
    cov_step: a callable that computes latent state and observation
      covariance at time `t`, given latent covariance at time `t-1`.
  """

  def cov_step(previous_covs, t):
    """Single step of prior covariance recursion."""
    previous_latent_cov, _ = previous_covs

    latent_cov = _propagate_cov(
        previous_latent_cov,
        get_transition_matrix_for_timestep(t - 1),
        get_transition_noise_for_timestep(t - 1))
    observation_cov = _propagate_cov(
        latent_cov,
        get_observation_matrix_for_timestep(t),
        get_observation_noise_for_timestep(t))

    return (latent_cov, observation_cov)

  return cov_step

Build a callable for one step of Kalman covariance recursion .


def build_kalman_sample_step(get_transition_matrix_for_timestep,
                             get_transition_noise_for_timestep,
                             get_observation_matrix_for_timestep,
                             get_observation_noise_for_timestep,
                             full_sample_and_batch_shape,
                             stream,
                             validate_args=False):
  """Build a callable for one step of Kalman sampling recursion.

  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]`.
    full_sample_and_batch_shape: Desired sample and batch shape of the
      returned samples, concatenated in a single `Tensor`.
    stream: `tfd.SeedStream` instance used to generate a
      sequence of random seeds.
    validate_args: if True, perform error checking at runtime.

  Returns:
    sample_step: a callable that samples the latent state and
      observation at time `t`, given latent state at time `t-1`.
  """

  def sample_step(sampled_prev, t):
    """Sample values for a single timestep."""
    latent_prev, _ = sampled_prev

    transition_matrix = get_transition_matrix_for_timestep(t - 1)
    transition_noise = get_transition_noise_for_timestep(t - 1)

    latent_pred = transition_matrix.matmul(latent_prev)
    latent_sampled = latent_pred + transition_noise.sample(
        sample_shape=_augment_sample_shape(
            transition_noise,
            full_sample_and_batch_shape,
            validate_args),
        seed=stream())[..., tf.newaxis]

    observation_matrix = get_observation_matrix_for_timestep(t)
    observation_noise = get_observation_noise_for_timestep(t)

    observation_pred = observation_matrix.matmul(latent_sampled)
    observation_sampled = observation_pred + observation_noise.sample(
        sample_shape=_augment_sample_shape(
            observation_noise,
            full_sample_and_batch_shape,
            validate_args),
        seed=stream())[..., tf.newaxis]

    return (latent_sampled, observation_sampled)

  return sample_step

def build_kalman_sample_step ( get_transition_matrix_for_timestep , get_transition_noise_for_timestep , get_observation_matrix_for_timestep , get_observation_noise_for_timestep , full_sample_and_batch_shape , stream , validate_args = False ) : def sample_step ( sampled_prev , t ) : latent_prev , _ = sampled_prev transition_matrix = get_transition_matrix_for_timestep ( t - 1 ) transition_noise = get_transition_noise_for_timestep ( t - 1 ) latent_pred = transition_matrix . matmul ( latent_prev ) latent_sampled = latent_pred + transition_noise . sample ( sample_shape = _augment_sample_shape ( transition_noise , full_sample_and_batch_shape , validate_args ) , seed = stream ( ) ) [ ... , tf . newaxis ] observation_matrix = get_observation_matrix_for_timestep ( t ) observation_noise = get_observation_noise_for_timestep ( t ) observation_pred = observation_matrix . matmul ( latent_sampled ) observation_sampled = observation_pred + observation_noise . sample ( sample_shape = _augment_sample_shape ( observation_noise , full_sample_and_batch_shape , validate_args ) , seed = stream ( ) ) [ ... , tf . newaxis ] return ( latent_sampled , observation_sampled ) return sample_step

Build a callable for one step of Kalman sampling recursion .


def _propagate_mean(mean, linop, dist):
  """Propagate a mean through linear Gaussian transformation."""
  return linop.matmul(mean) + dist.mean()[..., tf.newaxis]

def _propagate_mean ( mean , linop , dist ) : return linop . matmul ( mean ) + dist . mean ( ) [ ... , tf . newaxis ]

Propagate a mean through linear Gaussian transformation .


def _propagate_cov(cov, linop, dist):
  """Propagate covariance through linear Gaussian transformation."""
  # For linop A and input cov P, returns `A P A' + dist.cov()`
  return linop.matmul(linop.matmul(cov), adjoint_arg=True) + dist.covariance()

def _propagate_cov ( cov , linop , dist ) : return linop . matmul ( linop . matmul ( cov ) , adjoint_arg = True ) + dist . covariance ( )

Propagate covariance through linear Gaussian transformation .


def backward_smoothing_pass(self,
                              filtered_means,
                              filtered_covs,
                              predicted_means,
                              predicted_covs):
    """Run the backward pass in Kalman smoother.

    The backward smoothing is using Rauch, Tung and Striebel smoother as
    as discussed in section 18.3.2 of Kevin P. Murphy, 2012, Machine Learning:
    A Probabilistic Perspective, The MIT Press. The inputs are returned by
    `forward_filter` function.

    Args:
      filtered_means: Means of the per-timestep filtered marginal
        distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`.
      filtered_covs: Covariances of the per-timestep filtered marginal
        distributions p(z_t | x_{:t}), as a Tensor of shape
        `batch_shape + [num_timesteps, latent_size, latent_size]`.
      predicted_means: Means of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, latent_size]`.
      predicted_covs: Covariances of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `batch_shape + [num_timesteps, latent_size,
         latent_size]`.

    Returns:
      posterior_means: Means of the smoothed marginal distributions
        p(z_t | x_{1:T}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`,
        which is of the same shape as filtered_means.
      posterior_covs: Covariances of the smoothed marginal distributions
        p(z_t | x_{1:T}), as a Tensor of shape
        `batch_shape + [num_timesteps, latent_size, latent_size]`.
        which is of the same shape as filtered_covs.
    """
    with tf.name_scope("backward_pass"):
      filtered_means = tf.convert_to_tensor(
          value=filtered_means, name="filtered_means")
      filtered_covs = tf.convert_to_tensor(
          value=filtered_covs, name="filtered_covs")
      predicted_means = tf.convert_to_tensor(
          value=predicted_means, name="predicted_means")
      predicted_covs = tf.convert_to_tensor(
          value=predicted_covs, name="predicted_covs")

      # To scan over time dimension, we need to move 'num_timesteps' from the
      # event shape to the initial dimension of the tensor.
      filtered_means = distribution_util.move_dimension(filtered_means, -2, 0)
      filtered_covs = distribution_util.move_dimension(filtered_covs, -3, 0)
      predicted_means = distribution_util.move_dimension(predicted_means, -2, 0)
      predicted_covs = distribution_util.move_dimension(predicted_covs, -3, 0)

      # The means are assumed to be vectors. Adding a dummy index to
      # ensure the `matmul` op working smoothly.
      filtered_means = filtered_means[..., tf.newaxis]
      predicted_means = predicted_means[..., tf.newaxis]

      initial_backward_mean = predicted_means[-1, ...]
      initial_backward_cov = predicted_covs[-1, ...]

      num_timesteps = tf.shape(input=filtered_means)[0]
      initial_state = BackwardPassState(
          backward_mean=initial_backward_mean,
          backward_cov=initial_backward_cov,
          timestep=self.initial_step + num_timesteps - 1)

      update_step_fn = build_backward_pass_step(
          self.get_transition_matrix_for_timestep)

      # For backward pass, it scans the `elems` from last to first.
      posterior_states = tf.scan(update_step_fn,
                                 elems=(filtered_means,
                                        filtered_covs,
                                        predicted_means,
                                        predicted_covs),
                                 initializer=initial_state,
                                 reverse=True)

      # Move the time dimension back into the event shape.
      posterior_means = distribution_util.move_dimension(
          posterior_states.backward_mean[..., 0], 0, -2)
      posterior_covs = distribution_util.move_dimension(
          posterior_states.backward_cov, 0, -3)

      return (posterior_means, posterior_covs)

Run the backward pass in Kalman smoother .


def _joint_sample_n(self, n, seed=None):
    """Draw a joint sample from the prior over latents and observations."""

    with tf.name_scope("sample_n_joint"):
      stream = seed_stream.SeedStream(
          seed, salt="LinearGaussianStateSpaceModel_sample_n_joint")

      sample_and_batch_shape = distribution_util.prefer_static_value(
          tf.concat([[n], self.batch_shape_tensor()],
                    axis=0))

      # Sample the initial timestep from the prior.  Since we want
      # this sample to have full batch shape (not just the batch shape
      # of the self.initial_state_prior object which might in general be
      # smaller), we augment the sample shape to include whatever
      # extra batch dimensions are required.
      with tf.control_dependencies(self.runtime_assertions):
        initial_latent = self.initial_state_prior.sample(
            sample_shape=_augment_sample_shape(
                self.initial_state_prior,
                sample_and_batch_shape,
                self.validate_args),
            seed=stream())

        # Add a dummy dimension so that matmul() does matrix-vector
        # multiplication.
        initial_latent = initial_latent[..., tf.newaxis]

      initial_observation_matrix = (
          self.get_observation_matrix_for_timestep(self.initial_step))
      initial_observation_noise = (
          self.get_observation_noise_for_timestep(self.initial_step))

      initial_observation_pred = initial_observation_matrix.matmul(
          initial_latent)
      initial_observation = (initial_observation_pred +
                             initial_observation_noise.sample(
                                 sample_shape=_augment_sample_shape(
                                     initial_observation_noise,
                                     sample_and_batch_shape,
                                     self.validate_args),
                                 seed=stream())[..., tf.newaxis])

      sample_step = build_kalman_sample_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep,
          full_sample_and_batch_shape=sample_and_batch_shape,
          stream=stream,
          validate_args=self.validate_args)

      # Scan over all timesteps to sample latents and observations.
      (latents, observations) = tf.scan(
          sample_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent, initial_observation))

      # Combine the initial sampled timestep with the remaining timesteps.
      latents = tf.concat([initial_latent[tf.newaxis, ...],
                           latents], axis=0)
      observations = tf.concat([initial_observation[tf.newaxis, ...],
                                observations], axis=0)

      # Put dimensions back in order. The samples we've computed are
      # ordered by timestep, with shape `[num_timesteps, num_samples,
      # batch_shape, size, 1]` where `size` represents `latent_size`
      # or `observation_size` respectively. But timesteps are really
      # part of each probabilistic event, so we need to return a Tensor
      # of shape `[num_samples, batch_shape, num_timesteps, size]`.
      latents = tf.squeeze(latents, -1)
      latents = distribution_util.move_dimension(latents, 0, -2)
      observations = tf.squeeze(observations, -1)
      observations = distribution_util.move_dimension(observations, 0, -2)

    return latents, observations

Draw a joint sample from the prior over latents and observations .


def posterior_marginals(self, x, mask=None):
    """Run a Kalman smoother to return posterior mean and cov.

    Note that the returned values `smoothed_means` depend on the observed
    time series `x`, while the `smoothed_covs` are independent
    of the observed series; i.e., they depend only on the model itself.
    This means that the mean values have shape `concat([sample_shape(x),
    batch_shape, [num_timesteps, {latent/observation}_size]])`,
    while the covariances have shape `concat[(batch_shape, [num_timesteps,
    {latent/observation}_size, {latent/observation}_size]])`, which
    does not depend on the sample shape.

    This function only performs smoothing. If the user wants the
    intermediate values, which are returned by filtering pass `forward_filter`,
    one could get it by:
    ```
    (log_likelihoods,
     filtered_means, filtered_covs,
     predicted_means, predicted_covs,
     observation_means, observation_covs) = model.forward_filter(x)
    smoothed_means, smoothed_covs = model.backward_smoothing_pass(x)
    ```
    where `x` is an observation sequence.

    Args:
      x: a float-type `Tensor` with rightmost dimensions
        `[num_timesteps, observation_size]` matching
        `self.event_shape`. Additional dimensions must match or be
        broadcastable to `self.batch_shape`; any further dimensions
        are interpreted as a sample shape.
      mask: optional bool-type `Tensor` with rightmost dimension
        `[num_timesteps]`; `True` values specify that the value of `x`
        at that timestep is masked, i.e., not conditioned on. Additional
        dimensions must match or be broadcastable to `self.batch_shape`; any
        further dimensions must match or be broadcastable to the sample
        shape of `x`.
        Default value: `None`.

    Returns:
      smoothed_means: Means of the per-timestep smoothed
         distributions over latent states, p(x_{t} | x_{:T}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, observation_size]`.
      smoothed_covs: Covariances of the per-timestep smoothed
         distributions over latent states, p(x_{t} | x_{:T}), as a
         Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
         observation_size, observation_size]`. Note that the covariances depend
         only on the model and the mask, not on the data, so this may have fewer
         dimensions than `filtered_means`.
    """

    with tf.name_scope("smooth"):
      x = tf.convert_to_tensor(value=x, name="x")
      (_, filtered_means, filtered_covs,
       predicted_means, predicted_covs, _, _) = self.forward_filter(
           x, mask=mask)

      (smoothed_means, smoothed_covs) = self.backward_smoothing_pass(
          filtered_means, filtered_covs,
          predicted_means, predicted_covs)

      return (smoothed_means, smoothed_covs)

def posterior_marginals ( self , x , mask = None ) : with tf . name_scope ( "smooth" ) : x = tf . convert_to_tensor ( value = x , name = "x" ) ( _ , filtered_means , filtered_covs , predicted_means , predicted_covs , _ , _ ) = self . forward_filter ( x , mask = mask ) ( smoothed_means , smoothed_covs ) = self . backward_smoothing_pass ( filtered_means , filtered_covs , predicted_means , predicted_covs ) return ( smoothed_means , smoothed_covs )

Run a Kalman smoother to return posterior mean and cov .


def _joint_mean(self):
    """Compute prior means for all variables via dynamic programming.

    Returns:
      latent_means: Prior means of latent states `z_t`, as a `Tensor`
        of shape `batch_shape + [num_timesteps, latent_size]`
      observation_means: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size]`
    """

    with tf.name_scope("mean_joint"):

      # The initial timestep is a special case, since we sample the
      # latent state from the prior rather than the transition model.

      with tf.control_dependencies(self.runtime_assertions):
        # Broadcast to ensure we represent the full batch shape.
        initial_latent_mean = _broadcast_to_shape(
            self.initial_state_prior.mean()[..., tf.newaxis],
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, 1]], axis=0))

      initial_observation_mean = _propagate_mean(
          initial_latent_mean,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      mean_step = build_kalman_mean_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_means, observation_means) = tf.scan(
          mean_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_mean, initial_observation_mean))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_means = tf.concat([initial_latent_mean[tf.newaxis, ...],
                                latent_means], axis=0)
      observation_means = tf.concat([initial_observation_mean[tf.newaxis, ...],
                                     observation_means], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, 1]`, where `size`
      # is the dimension of the latent or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size]`.
      latent_means = tf.squeeze(latent_means, -1)
      latent_means = distribution_util.move_dimension(latent_means, 0, -2)
      observation_means = tf.squeeze(observation_means, -1)
      observation_means = distribution_util.move_dimension(
          observation_means, 0, -2)

      return latent_means, observation_means

def _joint_mean ( self ) : with tf . name_scope ( "mean_joint" ) : with tf . control_dependencies ( self . runtime_assertions ) : initial_latent_mean = _broadcast_to_shape ( self . initial_state_prior . mean ( ) [ ... , tf . newaxis ] , tf . concat ( [ self . batch_shape_tensor ( ) , [ self . latent_size , 1 ] ] , axis = 0 ) ) initial_observation_mean = _propagate_mean ( initial_latent_mean , self . get_observation_matrix_for_timestep ( self . initial_step ) , self . get_observation_noise_for_timestep ( self . initial_step ) ) mean_step = build_kalman_mean_step ( self . get_transition_matrix_for_timestep , self . get_transition_noise_for_timestep , self . get_observation_matrix_for_timestep , self . get_observation_noise_for_timestep ) ( latent_means , observation_means ) = tf . scan ( mean_step , elems = tf . range ( self . initial_step + 1 , self . final_step ) , initializer = ( initial_latent_mean , initial_observation_mean ) ) latent_means = tf . concat ( [ initial_latent_mean [ tf . newaxis , ... ] , latent_means ] , axis = 0 ) observation_means = tf . concat ( [ initial_observation_mean [ tf . newaxis , ... ] , observation_means ] , axis = 0 ) latent_means = tf . squeeze ( latent_means , - 1 ) latent_means = distribution_util . move_dimension ( latent_means , 0 , - 2 ) observation_means = tf . squeeze ( observation_means , - 1 ) observation_means = distribution_util . move_dimension ( observation_means , 0 , - 2 ) return latent_means , observation_means

Compute prior means for all variables via dynamic programming .


def _joint_covariances(self):
    """Compute prior covariances for all variables via dynamic programming.

    Returns:
      latent_covs: Prior covariance matrices of latent states `z_t`, as
        a `Tensor` of shape `batch_shape + [num_timesteps,
        latent_size, latent_size]`
      observation_covs: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size, observation_size]`
    """

    with tf.name_scope("covariance_joint"):

      with tf.control_dependencies(self.runtime_assertions):
        initial_latent_cov = _broadcast_to_shape(
            self.initial_state_prior.covariance(),
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, self.latent_size]], axis=0))

      initial_observation_cov = _propagate_cov(
          initial_latent_cov,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      cov_step = build_kalman_cov_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_covs, observation_covs) = tf.scan(
          cov_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_cov, initial_observation_cov))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_covs = tf.concat([initial_latent_cov[tf.newaxis, ...],
                               latent_covs], axis=0)
      observation_covs = tf.concat([initial_observation_cov[tf.newaxis, ...],
                                    observation_covs], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, size]`, where `size`
      # is the dimension of the state or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size, size]`.
      latent_covs = distribution_util.move_dimension(latent_covs, 0, -3)
      observation_covs = distribution_util.move_dimension(
          observation_covs, 0, -3)
      return latent_covs, observation_covs

def _joint_covariances ( self ) : with tf . name_scope ( "covariance_joint" ) : with tf . control_dependencies ( self . runtime_assertions ) : initial_latent_cov = _broadcast_to_shape ( self . initial_state_prior . covariance ( ) , tf . concat ( [ self . batch_shape_tensor ( ) , [ self . latent_size , self . latent_size ] ] , axis = 0 ) ) initial_observation_cov = _propagate_cov ( initial_latent_cov , self . get_observation_matrix_for_timestep ( self . initial_step ) , self . get_observation_noise_for_timestep ( self . initial_step ) ) cov_step = build_kalman_cov_step ( self . get_transition_matrix_for_timestep , self . get_transition_noise_for_timestep , self . get_observation_matrix_for_timestep , self . get_observation_noise_for_timestep ) ( latent_covs , observation_covs ) = tf . scan ( cov_step , elems = tf . range ( self . initial_step + 1 , self . final_step ) , initializer = ( initial_latent_cov , initial_observation_cov ) ) latent_covs = tf . concat ( [ initial_latent_cov [ tf . newaxis , ... ] , latent_covs ] , axis = 0 ) observation_covs = tf . concat ( [ initial_observation_cov [ tf . newaxis , ... ] , observation_covs ] , axis = 0 ) latent_covs = distribution_util . move_dimension ( latent_covs , 0 , - 3 ) observation_covs = distribution_util . move_dimension ( observation_covs , 0 , - 3 ) return latent_covs , observation_covs

Compute prior covariances for all variables via dynamic programming .


def latents_to_observations(self, latent_means, latent_covs):
    """Push latent means and covariances forward through the observation model.

    Args:
      latent_means: float `Tensor` of shape `[..., num_timesteps, latent_size]`
      latent_covs: float `Tensor` of shape
        `[..., num_timesteps, latent_size, latent_size]`.

    Returns:
      observation_means: float `Tensor` of shape
        `[..., num_timesteps, observation_size]`
      observation_covs: float `Tensor` of shape
        `[..., num_timesteps, observation_size, observation_size]`
    """

    with tf.name_scope("latents_to_observations"):

      pushforward_latents_step = build_pushforward_latents_step(
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      latent_means = distribution_util.move_dimension(
          latent_means, source_idx=-2, dest_idx=0)
      latent_means = latent_means[..., tf.newaxis]  # Make matmul happy.
      latent_covs = distribution_util.move_dimension(
          latent_covs, source_idx=-3, dest_idx=0)

      (initial_observation_mean,
       initial_observation_cov) = pushforward_latents_step(
           _=None,  # Loop body ignores previous observations.
           latent_t_mean_cov=(self.initial_step,
                              latent_means[self.initial_step],
                              latent_covs[self.initial_step]))

      # TODO(davmre) this loop is embarassingly parallel; replace with `pfor`.
      timesteps = tf.range(self.initial_step,
                           self.initial_step + self.num_timesteps)
      observation_means, observation_covs = tf.scan(
          pushforward_latents_step,
          elems=(timesteps, latent_means, latent_covs),
          initializer=(initial_observation_mean, initial_observation_cov),
          parallel_iterations=10000)

      observation_means = distribution_util.move_dimension(
          observation_means[..., 0], source_idx=0, dest_idx=-2)
      observation_covs = distribution_util.move_dimension(
          observation_covs, source_idx=0, dest_idx=-3)

      return observation_means, observation_covs

Push latent means and covariances forward through the observation model .


def _log_normalization(self):
    """Computes the log-normalizer of the distribution."""
    event_dim = tf.compat.dimension_value(self.event_shape[0])
    if event_dim is None:
      raise ValueError('vMF _log_normalizer currently only supports '
                       'statically known event shape')
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_lognorm = ((event_dim / 2 - 1) * tf.math.log(safe_conc) -
                    (event_dim / 2) * np.log(2 * np.pi) -
                    tf.math.log(_bessel_ive(event_dim / 2 - 1, safe_conc)) -
                    tf.abs(safe_conc))
    log_nsphere_surface_area = (
        np.log(2.) + (event_dim / 2) * np.log(np.pi) -
        tf.math.lgamma(tf.cast(event_dim / 2, self.dtype)))
    return tf.where(self.concentration > 0,
                    -safe_lognorm,
                    log_nsphere_surface_area * tf.ones_like(safe_lognorm))

Computes the log - normalizer of the distribution .


def _mode(self):
    """The mode of the von Mises-Fisher distribution is the mean direction."""
    return (self.mean_direction +
            tf.zeros_like(self.concentration)[..., tf.newaxis])

def _mode ( self ) : return ( self . mean_direction + tf . zeros_like ( self . concentration ) [ ... , tf . newaxis ] )

The mode of the von Mises - Fisher distribution is the mean direction .


def _rotate(self, samples):
    """Applies a Householder rotation to `samples`."""
    event_dim = (
        tf.compat.dimension_value(self.event_shape[0]) or
        self._event_shape_tensor()[0])
    basis = tf.concat([[1.], tf.zeros([event_dim - 1], dtype=self.dtype)],
                      axis=0),
    u = tf.nn.l2_normalize(basis - self.mean_direction, axis=-1)
    return samples - 2 * tf.reduce_sum(
        input_tensor=samples * u, axis=-1, keepdims=True) * u

def _rotate ( self , samples ) : event_dim = ( tf . compat . dimension_value ( self . event_shape [ 0 ] ) or self . _event_shape_tensor ( ) [ 0 ] ) basis = tf . concat ( [ [ 1. ] , tf . zeros ( [ event_dim - 1 ] , dtype = self . dtype ) ] , axis = 0 ) , u = tf . nn . l2_normalize ( basis - self . mean_direction , axis = - 1 ) return samples - 2 * tf . reduce_sum ( input_tensor = samples * u , axis = - 1 , keepdims = True ) * u

Applies a Householder rotation to samples .


def _sample_3d(self, n, seed=None):
    """Specialized inversion sampler for 3D."""
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]

def
_sample_3d
(
self
,
n
,
seed
=
)
:
seed
=
seed_stream
.
SeedStream
(
seed
,
salt
=
'von_mises_fisher_3d'
)
u_shape
=
tf
.
concat
(
[
[
n
]
,
self
.
_batch_shape_tensor
(
)
]
,
axis
=
0
)
z
=
tf
.
random
.
uniform
(
u_shape
,
seed
=
seed
(
)
,
dtype
=
self
.
dtype
)
safe_conc
=
tf
.
where
(
self
.
concentration
>
0
,
self
.
concentration
,
tf
.
ones_like
(
self
.
concentration
)
)
safe_z
=
tf
.
where
(
z
>
0
,
z
,
tf
.
ones_like
(
z
)
)
safe_u
=
1
+
tf
.
reduce_logsumexp
(
input_tensor
=
[
tf
.
math
.
log
(
safe_z
)
,
tf
.
math
.
log1p
(
safe_z
)
2
*
safe_conc
]
,
axis
=
0
)
/
safe_conc
u
=
tf
.
where
(
self
.
concentration
>
tf
.
zeros_like
(
safe_u
)
,
safe_u
,
2
*
z
1
)
u
=
tf
.
where
(
tf
.
equal
(
z
,
0
)
,
tf
.
ones_like
(
u
)
,
u
)
if
not
self
.
_allow_nan_stats
:
u
=
tf
.
debugging
.
check_numerics
(
u
,
'u in _sample_3d'
)
return
u
[
...
,
tf
.
newaxis
]


def _copy_fn(fn):
  """Create a deep copy of fn.

  Args:
    fn: a callable

  Returns:
    A `FunctionType`: a deep copy of fn.

  Raises:
    TypeError: if `fn` is not a callable.
  """
  if not callable(fn):
    raise TypeError("fn is not callable: {}".format(fn))
  # The blessed way to copy a function. copy.deepcopy fails to create a
  # non-reference copy. Since:
  #   types.FunctionType == type(lambda: None),
  # and the docstring for the function type states:
  #
  #   function(code, globals[, name[, argdefs[, closure]]])
  #
  #   Create a function object from a code object and a dictionary.
  #   ...
  #
  # Here we can use this to create a new function with the old function's
  # code, globals, closure, etc.
  return types.FunctionType(
      code=fn.__code__, globals=fn.__globals__,
      name=fn.__name__, argdefs=fn.__defaults__,
      closure=fn.__closure__)

def
_copy_fn
(
fn
)
:
if
not
callable
(
fn
)
:
raise
TypeError
(
"fn is not callable: {}"
.
format
(
fn
)
)
return
types
.
FunctionType
(
code
=
fn
.
__code__
,
globals
=
fn
.
__globals__
,
name
=
fn
.
__name__
,
argdefs
=
fn
.
__defaults__
,
closure
=
fn
.
__closure__
)


def _remove_dict_keys_with_value(dict_, val):
  """Removes `dict` keys which have have `self` as value."""
  return {k: v for k, v in dict_.items() if v is not val}

def _remove_dict_keys_with_value ( dict_ , val ) : return { k : v for k , v in dict_ . items ( ) if v is not val }

Removes dict keys which have have self as value .


def _recursively_replace_dict_for_pretty_dict(x):
  """Recursively replace `dict`s with `_PrettyDict`."""
  # We use "PrettyDict" because collections.OrderedDict repr/str has the word
  # "OrderedDict" in it. We only want to print "OrderedDict" if in fact the
  # input really is an OrderedDict.
  if isinstance(x, dict):
    return _PrettyDict({
        k: _recursively_replace_dict_for_pretty_dict(v)
        for k, v in x.items()})
  if (isinstance(x, collections.Sequence) and
      not isinstance(x, six.string_types)):
    args = (_recursively_replace_dict_for_pretty_dict(x_) for x_ in x)
    is_named_tuple = (isinstance(x, tuple) and
                      hasattr(x, "_asdict") and
                      hasattr(x, "_fields"))
    return type(x)(*args) if is_named_tuple else type(x)(args)
  if isinstance(x, collections.Mapping):
    return type(x)(**{k: _recursively_replace_dict_for_pretty_dict(v)
                      for k, v in x.items()})
  return x

Recursively replace dict s with _PrettyDict .


def _get_samples(dist, z, n, seed):
  """Check args and return samples."""
  with tf.compat.v1.name_scope('get_samples', values=[z, n]):
    if (n is None) == (z is None):
      raise ValueError(
          'Must specify exactly one of arguments "n" and "z".  Found: '
          'n = %s, z = %s' % (n, z))
    if n is not None:
      return dist.sample(n, seed=seed)
    else:
      return tf.convert_to_tensor(value=z, name='z')


def is_namedtuple_like(x):
  """Helper which returns `True` if input is `collections.namedtuple`-like."""
  try:
    for fn in x._fields:
      _ = getattr(x, fn)
    return True
  except AttributeError:
    return False

def is_namedtuple_like ( x ) : try : for fn in x . _fields : _ = getattr ( x , fn ) return True except AttributeError : return False

Helper which returns True if input is collections . namedtuple - like .


def _choose_base_case(is_accepted,
                      accepted,
                      rejected,
                      name=None):
  """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
  def _expand_is_accepted_like(x):
    """Helper to expand `is_accepted` like the shape of some input arg."""
    with tf.compat.v1.name_scope('expand_is_accepted_like'):
      expand_shape = tf.concat([
          tf.shape(input=is_accepted),
          tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),
      ],
                               axis=0)
      multiples = tf.concat([
          tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
          tf.shape(input=x)[tf.rank(is_accepted):],
      ],
                            axis=0)
      m = tf.tile(tf.reshape(is_accepted, expand_shape),
                  multiples)
      m.set_shape(m.shape.merge_with(x.shape))
      return m
  def _where(accepted, rejected):
    if accepted is rejected:
      return accepted
    accepted = tf.convert_to_tensor(value=accepted, name='accepted')
    rejected = tf.convert_to_tensor(value=rejected, name='rejected')
    r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)
    r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))
    return r

  with tf.compat.v1.name_scope(
      name, 'choose', values=[is_accepted, accepted, rejected]):
    if not is_list_like(accepted):
      return _where(accepted, rejected)
    return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)
             else _where(a, r))
            for a, r in zip(accepted, rejected)]

def _choose_base_case ( is_accepted , accepted , rejected , name = None ) : def _expand_is_accepted_like ( x ) : with tf . compat . v1 . name_scope ( 'expand_is_accepted_like' ) : expand_shape = tf . concat ( [ tf . shape ( input = is_accepted ) , tf . ones ( [ tf . rank ( x ) - tf . rank ( is_accepted ) ] , dtype = tf . int32 ) , ] , axis = 0 ) multiples = tf . concat ( [ tf . ones ( [ tf . rank ( is_accepted ) ] , dtype = tf . int32 ) , tf . shape ( input = x ) [ tf . rank ( is_accepted ) : ] , ] , axis = 0 ) m = tf . tile ( tf . reshape ( is_accepted , expand_shape ) , multiples ) m . set_shape ( m . shape . merge_with ( x . shape ) ) return m def _where ( accepted , rejected ) : if accepted is rejected : return accepted accepted = tf . convert_to_tensor ( value = accepted , name = 'accepted' ) rejected = tf . convert_to_tensor ( value = rejected , name = 'rejected' ) r = tf . where ( _expand_is_accepted_like ( accepted ) , accepted , rejected ) r . set_shape ( r . shape . merge_with ( accepted . shape . merge_with ( rejected . shape ) ) ) return r with tf . compat . v1 . name_scope ( name , 'choose' , values = [ is_accepted , accepted , rejected ] ) : if not is_list_like ( accepted ) : return _where ( accepted , rejected ) return [ ( choose ( is_accepted , a , r , name = name ) if is_namedtuple_like ( a ) else _where ( a , r ) ) for a , r in zip ( accepted , rejected ) ]

Helper to choose which expand_dims is_accepted and applies tf . where .


def choose(is_accepted, accepted, rejected, name=None):
  """Helper which expand_dims `is_accepted` then applies tf.where."""
  if not is_namedtuple_like(accepted):
    return _choose_base_case(is_accepted, accepted, rejected, name=name)
  if not isinstance(accepted, type(rejected)):
    raise TypeError('Type of `accepted` ({}) must be identical to '
                    'type of `rejected` ({})'.format(
                        type(accepted).__name__,
                        type(rejected).__name__))
  return type(accepted)(**dict(
      [(fn,
        choose(is_accepted,
               getattr(accepted, fn),
               getattr(rejected, fn),
               name=name))
       for fn in accepted._fields]))

Helper which expand_dims is_accepted then applies tf . where .


def safe_sum(x, alt_value=-np.inf, name=None):
  """Elementwise adds list members, replacing non-finite results with alt_value.

  Typically the `alt_value` is chosen so the `MetropolisHastings`
  `TransitionKernel` always rejects the proposal.

  Args:
    x: Python `list` of `Tensors` to elementwise add.
    alt_value: Python scalar used to replace any elementwise sums which would
      otherwise be non-finite.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "safe_sum").

  Returns:
    safe_sum: `Tensor` representing the elementwise sum of list of `Tensor`s
      `x` or `alt_value` where sums are non-finite.

  Raises:
    TypeError: if `x` is not list-like.
    ValueError: if `x` is empty.
  """
  with tf.compat.v1.name_scope(name, 'safe_sum', [x, alt_value]):
    if not is_list_like(x):
      raise TypeError('Expected list input.')
    if not x:
      raise ValueError('Input should not be empty.')
    in_shape = x[0].shape
    x = tf.stack(x, axis=-1)
    x = tf.reduce_sum(input_tensor=x, axis=-1)
    alt_value = np.array(alt_value, x.dtype.as_numpy_dtype)
    alt_fill = tf.fill(tf.shape(input=x), value=alt_value)
    x = tf.where(tf.math.is_finite(x), x, alt_fill)
    x.set_shape(x.shape.merge_with(in_shape))
    return x

def
safe_sum
(
x
,
alt_value
=
np
.
inf
,
name
=
)
:
with
tf
.
compat
.
v1
.
name_scope
(
name
,
'safe_sum'
,
[
x
,
alt_value
]
)
:
if
not
is_list_like
(
x
)
:
raise
TypeError
(
'Expected list input.'
)
if
not
x
:
raise
ValueError
(
'Input should not be empty.'
)
in_shape
=
x
[
0
]
.
shape
x
=
tf
.
stack
(
x
,
axis
=
1
)
x
=
tf
.
reduce_sum
(
input_tensor
=
x
,
axis
=
1
)
alt_value
=
np
.
array
(
alt_value
,
x
.
dtype
.
as_numpy_dtype
)
alt_fill
=
tf
.
fill
(
tf
.
shape
(
input
=
x
)
,
value
=
alt_value
)
x
=
tf
.
where
(
tf
.
math
.
is_finite
(
x
)
,
x
,
alt_fill
)
x
.
set_shape
(
x
.
shape
.
merge_with
(
in_shape
)
)
return
x

Elementwise adds list members replacing non - finite results with alt_value .


def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):
  """Helper to `maybe_call_fn_and_grads`."""
  with tf.compat.v1.name_scope(name, 'value_and_gradients',
                               [fn_arg_list, result, grads]):

    def _convert_to_tensor(x, name):
      ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(
          value=x_, name=name)
      return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)

    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    fn_arg_list = _convert_to_tensor(fn_arg_list, 'fn_arg')

    if result is None:
      result = fn(*fn_arg_list)
      if grads is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        fn_arg_list = [0 + x for x in fn_arg_list]

    result = _convert_to_tensor(result, 'fn_result')

    if grads is not None:
      grads = _convert_to_tensor(grads, 'fn_grad')
      return result, grads

    if is_list_like(result) and len(result) == len(fn_arg_list):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(i):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+1:]))
      grads = [
          tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[1]
          for i in range(len(result))
      ]
    else:
      _, grads = tfp_math_value_and_gradients(fn, fn_arg_list)

    return result, grads

def _value_and_gradients ( fn , fn_arg_list , result = None , grads = None , name = None ) : with tf . compat . v1 . name_scope ( name , 'value_and_gradients' , [ fn_arg_list , result , grads ] ) : def _convert_to_tensor ( x , name ) : ctt = lambda x_ : x_ if x_ is None else tf . convert_to_tensor ( value = x_ , name = name ) return [ ctt ( x_ ) for x_ in x ] if is_list_like ( x ) else ctt ( x ) fn_arg_list = ( list ( fn_arg_list ) if is_list_like ( fn_arg_list ) else [ fn_arg_list ] ) fn_arg_list = _convert_to_tensor ( fn_arg_list , 'fn_arg' ) if result is None : result = fn ( * fn_arg_list ) if grads is None and tf . executing_eagerly ( ) : fn_arg_list = [ 0 + x for x in fn_arg_list ] result = _convert_to_tensor ( result , 'fn_result' ) if grads is not None : grads = _convert_to_tensor ( grads , 'fn_grad' ) return result , grads if is_list_like ( result ) and len ( result ) == len ( fn_arg_list ) : def fn_slice ( i ) : return lambda x : fn ( * ( fn_arg_list [ : i ] + [ x ] + fn_arg_list [ i + 1 : ] ) ) grads = [ tfp_math_value_and_gradients ( fn_slice ( i ) , fn_arg_list [ i ] ) [ 1 ] for i in range ( len ( result ) ) ] else : _ , grads = tfp_math_value_and_gradients ( fn , fn_arg_list ) return result , grads


def maybe_call_fn_and_grads(fn,
                            fn_arg_list,
                            result=None,
                            grads=None,
                            check_non_none_grads=True,
                            name=None):
  """Calls `fn` and computes the gradient of the result wrt `args_list`."""
  with tf.compat.v1.name_scope(name, 'maybe_call_fn_and_grads',
                               [fn_arg_list, result, grads]):
    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    result, grads = _value_and_gradients(fn, fn_arg_list, result, grads)
    if not all(r.dtype.is_floating
               for r in (result if is_list_like(result) else [result])):  # pylint: disable=superfluous-parens
      raise TypeError('Function result must be a `Tensor` with `float` '
                      '`dtype`.')
    if len(fn_arg_list) != len(grads):
      raise ValueError('Function args must be in one-to-one correspondence '
                       'with grads.')
    if check_non_none_grads and any(g is None for g in grads):
      raise ValueError('Encountered `None` gradient.\n'
                       '  fn_arg_list: {}\n'
                       '  grads: {}'.format(fn_arg_list, grads))
    return result, grads

Calls fn and computes the gradient of the result wrt args_list .


def smart_for_loop(loop_num_iter, body_fn, initial_loop_vars,
                   parallel_iterations=10, name=None):
  """Construct a for loop, preferring a python loop if `n` is staticaly known.

  Given `loop_num_iter` and `body_fn`, return an op corresponding to executing
  `body_fn` `loop_num_iter` times, feeding previous outputs of `body_fn` into
  the next iteration.

  If `loop_num_iter` is statically known, the op is constructed via python for
  loop, and otherwise a `tf.while_loop` is used.

  Args:
    loop_num_iter: `Integer` `Tensor` representing the number of loop
      iterations.
    body_fn: Callable to be executed `loop_num_iter` times.
    initial_loop_vars: Listlike object of `Tensors` to be passed in to
      `body_fn`'s first execution.
    parallel_iterations: The number of iterations allowed to run in parallel.
      It must be a positive integer. See `tf.while_loop` for more details.
      Default value: `10`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "smart_for_loop").
  Returns:
    result: `Tensor` representing applying `body_fn` iteratively `n` times.
  """
  with tf.compat.v1.name_scope(name, 'smart_for_loop',
                               [loop_num_iter, initial_loop_vars]):
    loop_num_iter_ = tf.get_static_value(loop_num_iter)
    if (loop_num_iter_ is None or tf.executing_eagerly() or
        control_flow_util.GraphOrParentsInXlaContext(
            tf.compat.v1.get_default_graph())):
      # Cast to int32 to run the comparison against i in host memory,
      # where while/LoopCond needs it.
      loop_num_iter = tf.cast(loop_num_iter, dtype=tf.int32)
      return tf.while_loop(
          cond=lambda i, *args: i < loop_num_iter,
          body=lambda i, *args: [i + 1] + list(body_fn(*args)),
          loop_vars=[np.int32(0)] + initial_loop_vars,
          parallel_iterations=parallel_iterations
      )[1:]
    result = initial_loop_vars
    for _ in range(loop_num_iter_):
      result = body_fn(*result)
    return result

def smart_for_loop ( loop_num_iter , body_fn , initial_loop_vars , parallel_iterations = 10 , name = None ) : with tf . compat . v1 . name_scope ( name , 'smart_for_loop' , [ loop_num_iter , initial_loop_vars ] ) : loop_num_iter_ = tf . get_static_value ( loop_num_iter ) if ( loop_num_iter_ is None or tf . executing_eagerly ( ) or control_flow_util . GraphOrParentsInXlaContext ( tf . compat . v1 . get_default_graph ( ) ) ) : loop_num_iter = tf . cast ( loop_num_iter , dtype = tf . int32 ) return tf . while_loop ( cond = lambda i , * args : i < loop_num_iter , body = lambda i , * args : [ i + 1 ] + list ( body_fn ( * args ) ) , loop_vars = [ np . int32 ( 0 ) ] + initial_loop_vars , parallel_iterations = parallel_iterations ) [ 1 : ] result = initial_loop_vars for _ in range ( loop_num_iter_ ) : result = body_fn ( * result ) return result

Construct a for loop preferring a python loop if n is staticaly known .


def trace_scan(loop_fn,
               initial_state,
               elems,
               trace_fn,
               parallel_iterations=10,
               name=None):
  """A simplified version of `tf.scan` that has configurable tracing.

  This function repeatedly calls `loop_fn(state, elem)`, where `state` is the
  `initial_state` during the first iteration, and the return value of `loop_fn`
  for every iteration thereafter. `elem` is a slice of `elements` along the
  first dimension, accessed in order. Additionally, it calls `trace_fn` on the
  return value of `loop_fn`. The `Tensor`s in return values of `trace_fn` are
  stacked and returned from this function, such that the first dimension of
  those `Tensor`s matches the size of `elems`.

  Args:
    loop_fn: A callable that takes in a `Tensor` or a nested collection of
      `Tensor`s with the same structure as `initial_state`, a slice of `elems`
      and returns the same structure as `initial_state`.
    initial_state: A `Tensor` or a nested collection of `Tensor`s passed to
      `loop_fn` in the first iteration.
    elems: A `Tensor` that is split along the first dimension and each element
      of which is passed to `loop_fn`.
    trace_fn: A callable that takes in the return value of `loop_fn` and returns
      a `Tensor` or a nested collection of `Tensor`s.
    parallel_iterations: Passed to the internal `tf.while_loop`.
    name: Name scope used in this function. Default: 'trace_scan'.

  Returns:
    final_state: The final return value of `loop_fn`.
    trace: The same structure as the return value of `trace_fn`, but with each
      `Tensor` being a stack of the corresponding `Tensors` in the return value
      of `trace_fn` for each slice of `elems`.
  """
  with tf.compat.v1.name_scope(
      name, 'trace_scan', [initial_state, elems]), tf.compat.v1.variable_scope(
          tf.compat.v1.get_variable_scope()) as vs:
    if vs.caching_device is None and not tf.executing_eagerly():
      vs.set_caching_device(lambda op: op.device)

    initial_state = tf.nest.map_structure(
        lambda x: tf.convert_to_tensor(value=x, name='initial_state'),
        initial_state)
    elems = tf.convert_to_tensor(value=elems, name='elems')

    static_length = elems.shape[0]
    if tf.compat.dimension_value(static_length) is None:
      length = tf.shape(input=elems)[0]
    else:
      length = tf.convert_to_tensor(
          value=static_length, dtype=tf.int32, name='length')

    # This is an TensorArray in part because of XLA, which had trouble with
    # non-statically known indices. I.e. elems[i] errored, but
    # elems_array.read(i) worked.
    elems_array = tf.TensorArray(
        elems.dtype, size=length, element_shape=elems.shape[1:])
    elems_array = elems_array.unstack(elems)

    trace_arrays = tf.nest.map_structure(
        lambda x: tf.TensorArray(x.dtype, size=length, element_shape=x.shape),
        trace_fn(initial_state))

    def _body(i, state, trace_arrays):
      state = loop_fn(state, elems_array.read(i))
      trace_arrays = tf.nest.pack_sequence_as(trace_arrays, [
          a.write(i, v) for a, v in zip(
              tf.nest.flatten(trace_arrays), tf.nest.flatten(trace_fn(state)))
      ])
      return i + 1, state, trace_arrays

    _, final_state, trace_arrays = tf.while_loop(
        cond=lambda i, *args: i < length,
        body=_body,
        loop_vars=(0, initial_state, trace_arrays),
        parallel_iterations=parallel_iterations)

    stacked_trace = tf.nest.map_structure(lambda x: x.stack(), trace_arrays)

    # Restore the static length if we know it.
    def _merge_static_length(x):
      x.set_shape(tf.TensorShape(static_length).concatenate(x.shape[1:]))
      return x

    stacked_trace = tf.nest.map_structure(_merge_static_length, stacked_trace)
    return final_state, stacked_trace

def trace_scan ( loop_fn , initial_state , elems , trace_fn , parallel_iterations = 10 , name = None ) : with tf . compat . v1 . name_scope ( name , 'trace_scan' , [ initial_state , elems ] ) , tf . compat . v1 . variable_scope ( tf . compat . v1 . get_variable_scope ( ) ) as vs : if vs . caching_device is None and not tf . executing_eagerly ( ) : vs . set_caching_device ( lambda op : op . device ) initial_state = tf . nest . map_structure ( lambda x : tf . convert_to_tensor ( value = x , name = 'initial_state' ) , initial_state ) elems = tf . convert_to_tensor ( value = elems , name = 'elems' ) static_length = elems . shape [ 0 ] if tf . compat . dimension_value ( static_length ) is None : length = tf . shape ( input = elems ) [ 0 ] else : length = tf . convert_to_tensor ( value = static_length , dtype = tf . int32 , name = 'length' ) elems_array = tf . TensorArray ( elems . dtype , size = length , element_shape = elems . shape [ 1 : ] ) elems_array = elems_array . unstack ( elems ) trace_arrays = tf . nest . map_structure ( lambda x : tf . TensorArray ( x . dtype , size = length , element_shape = x . shape ) , trace_fn ( initial_state ) ) def _body ( i , state , trace_arrays ) : state = loop_fn ( state , elems_array . read ( i ) ) trace_arrays = tf . nest . pack_sequence_as ( trace_arrays , [ a . write ( i , v ) for a , v in zip ( tf . nest . flatten ( trace_arrays ) , tf . nest . flatten ( trace_fn ( state ) ) ) ] ) return i + 1 , state , trace_arrays _ , final_state , trace_arrays = tf . while_loop ( cond = lambda i , * args : i < length , body = _body , loop_vars = ( 0 , initial_state , trace_arrays ) , parallel_iterations = parallel_iterations ) stacked_trace = tf . nest . map_structure ( lambda x : x . stack ( ) , trace_arrays ) def _merge_static_length ( x ) : x . set_shape ( tf . TensorShape ( static_length ) . concatenate ( x . shape [ 1 : ] ) ) return x stacked_trace = tf . nest . map_structure ( _merge_static_length , stacked_trace ) return final_state , stacked_trace

A simplified version of tf . scan that has configurable tracing .


def make_innermost_setter(setter):
  """Wraps a setter so it applies to the inner-most results in `kernel_results`.

  The wrapped setter unwraps `kernel_results` and applies `setter` to the first
  results without an `inner_results` attribute.

  Args:
    setter: A callable that takes the kernel results as well as some `*args` and
      `**kwargs` and returns a modified copy of those kernel results.

  Returns:
    new_setter: A wrapped `setter`.
  """

  @functools.wraps(setter)
  def _new_setter(kernel_results, *args, **kwargs):
    """Wrapped setter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    new_kernel_results = setter(kernel_results, *args, **kwargs)
    for outer_results in reversed(results_stack):
      new_kernel_results = outer_results._replace(
          inner_results=new_kernel_results)

    return new_kernel_results

  return _new_setter

def make_innermost_setter ( setter ) : @ functools . wraps ( setter ) def _new_setter ( kernel_results , * args , ** kwargs ) : results_stack = [ ] while hasattr ( kernel_results , 'inner_results' ) : results_stack . append ( kernel_results ) kernel_results = kernel_results . inner_results new_kernel_results = setter ( kernel_results , * args , ** kwargs ) for outer_results in reversed ( results_stack ) : new_kernel_results = outer_results . _replace ( inner_results = new_kernel_results ) return new_kernel_results return _new_setter

Wraps a setter so it applies to the inner - most results in kernel_results .


def make_innermost_getter(getter):
  """Wraps a getter so it applies to the inner-most results in `kernel_results`.

  The wrapped getter unwraps `kernel_results` and returns the return value of
  `getter` called with the first results without an `inner_results` attribute.

  Args:
    getter: A callable that takes Kernel results and returns some value.

  Returns:
    new_getter: A wrapped `getter`.
  """

  @functools.wraps(getter)
  def _new_getter(kernel_results, *args, **kwargs):
    """Wrapped getter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    return getter(kernel_results, *args, **kwargs)

  return _new_getter

def make_innermost_getter ( getter ) : @ functools . wraps ( getter ) def _new_getter ( kernel_results , * args , ** kwargs ) : results_stack = [ ] while hasattr ( kernel_results , 'inner_results' ) : results_stack . append ( kernel_results ) kernel_results = kernel_results . inner_results return getter ( kernel_results , * args , ** kwargs ) return _new_getter

Wraps a getter so it applies to the inner - most results in kernel_results .


def enable_store_parameters_in_results(kernel):
  """Enables the `store_parameters_in_results` parameter in a chain of kernels.

  This is a temporary utility for use during the transition period of the
  parameter storage methods.

  Args:
    kernel: A TransitionKernel.

  Returns:
    kernel: The same kernel, but recreated with `store_parameters_in_results`
        recursively set to `True` in its parameters and its inner kernels (as
        appropriate).
  """
  kernel_stack = []
  while hasattr(kernel, 'parameters') and 'inner_kernel' in kernel.parameters:
    kernel_stack.append(kernel)
    kernel = kernel.parameters['inner_kernel']

  def _recreate_kernel(kernel, parameters):
    new_parameters = kernel.parameters.copy()
    new_parameters.update(parameters)
    if 'store_parameters_in_results' in new_parameters:
      new_parameters['store_parameters_in_results'] = True
    with deprecation.silence():
      return type(kernel)(**new_parameters)

  if hasattr(kernel, 'parameters'):
    kernel = _recreate_kernel(kernel, {})

  for outer_kernel in reversed(kernel_stack):
    outer_kernel = _recreate_kernel(outer_kernel, {'inner_kernel': kernel})
    kernel = outer_kernel

  return kernel

def enable_store_parameters_in_results ( kernel ) : kernel_stack = [ ] while hasattr ( kernel , 'parameters' ) and 'inner_kernel' in kernel . parameters : kernel_stack . append ( kernel ) kernel = kernel . parameters [ 'inner_kernel' ] def _recreate_kernel ( kernel , parameters ) : new_parameters = kernel . parameters . copy ( ) new_parameters . update ( parameters ) if 'store_parameters_in_results' in new_parameters : new_parameters [ 'store_parameters_in_results' ] = True with deprecation . silence ( ) : return type ( kernel ) ( ** new_parameters ) if hasattr ( kernel , 'parameters' ) : kernel = _recreate_kernel ( kernel , { } ) for outer_kernel in reversed ( kernel_stack ) : outer_kernel = _recreate_kernel ( outer_kernel , { 'inner_kernel' : kernel } ) kernel = outer_kernel return kernel

Enables the store_parameters_in_results parameter in a chain of kernels .


def _replace_event_shape_in_shape_tensor(
    input_shape, event_shape_in, event_shape_out, validate_args):
  """Replaces the rightmost dims in a `Tensor` representing a shape.

  Args:
    input_shape: a rank-1 `Tensor` of integers
    event_shape_in: the event shape expected to be present in rightmost dims
      of `shape_in`.
    event_shape_out: the event shape with which to replace `event_shape_in` in
      the rightmost dims of `input_shape`.
    validate_args: Python `bool` indicating whether arguments should
      be checked for correctness.

  Returns:
    output_shape: A rank-1 integer `Tensor` with the same contents as
      `input_shape` except for the event dims, which are replaced with
      `event_shape_out`.
  """
  output_tensorshape, is_validated = _replace_event_shape_in_tensorshape(
      tensorshape_util.constant_value_as_shape(input_shape),
      event_shape_in,
      event_shape_out)

  # TODO(b/124240153): Remove map(tf.identity, deps) once tf.function
  # correctly supports control_dependencies.
  validation_dependencies = (
      map(tf.identity, (event_shape_in, event_shape_out))
      if validate_args else ())

  if (tensorshape_util.is_fully_defined(output_tensorshape) and
      (is_validated or not validate_args)):
    with tf.control_dependencies(validation_dependencies):
      output_shape = tf.convert_to_tensor(
          value=output_tensorshape, name='output_shape', dtype_hint=tf.int32)
    return output_shape, output_tensorshape

  with tf.control_dependencies(validation_dependencies):
    event_shape_in_ndims = (
        tf.size(input=event_shape_in)
        if tensorshape_util.num_elements(event_shape_in.shape) is None else
        tensorshape_util.num_elements(event_shape_in.shape))
    input_non_event_shape, input_event_shape = tf.split(
        input_shape, num_or_size_splits=[-1, event_shape_in_ndims])

  additional_assertions = []
  if is_validated:
    pass
  elif validate_args:
    # Check that `input_event_shape` and `event_shape_in` are compatible in the
    # sense that they have equal entries in any position that isn't a `-1` in
    # `event_shape_in`. Note that our validations at construction time ensure
    # there is at most one such entry in `event_shape_in`.
    mask = event_shape_in >= 0
    explicit_input_event_shape = tf.boolean_mask(
        tensor=input_event_shape, mask=mask)
    explicit_event_shape_in = tf.boolean_mask(
        tensor=event_shape_in, mask=mask)
    additional_assertions.append(
        assert_util.assert_equal(
            explicit_input_event_shape,
            explicit_event_shape_in,
            message='Input `event_shape` does not match `event_shape_in`.'))
    # We don't explicitly additionally verify
    # `tf.size(input_shape) > tf.size(event_shape_in)` since `tf.split`
    # already makes this assertion.

  with tf.control_dependencies(additional_assertions):
    output_shape = tf.concat([input_non_event_shape, event_shape_out], axis=0,
                             name='output_shape')

  return output_shape, output_tensorshape

Replaces the rightmost dims in a Tensor representing a shape .


def _replace_event_shape_in_tensorshape(
    input_tensorshape, event_shape_in, event_shape_out):
  """Replaces the event shape dims of a `TensorShape`.

  Args:
    input_tensorshape: a `TensorShape` instance in which to attempt replacing
      event shape.
    event_shape_in: `Tensor` shape representing the event shape expected to
      be present in (rightmost dims of) `tensorshape_in`. Must be compatible
      with the rightmost dims of `tensorshape_in`.
    event_shape_out: `Tensor` shape representing the new event shape, i.e.,
      the replacement of `event_shape_in`,

  Returns:
    output_tensorshape: `TensorShape` with the rightmost `event_shape_in`
      replaced by `event_shape_out`. Might be partially defined, i.e.,
      `TensorShape(None)`.
    is_validated: Python `bool` indicating static validation happened.

  Raises:
    ValueError: if we can determine the event shape portion of
      `tensorshape_in` as well as `event_shape_in` both statically, and they
      are not compatible. "Compatible" here means that they are identical on
      any dims that are not -1 in `event_shape_in`.
  """
  event_shape_in_ndims = tensorshape_util.num_elements(event_shape_in.shape)
  if tensorshape_util.rank(
      input_tensorshape) is None or event_shape_in_ndims is None:
    return tf.TensorShape(None), False  # Not is_validated.

  input_non_event_ndims = tensorshape_util.rank(
      input_tensorshape) - event_shape_in_ndims
  if input_non_event_ndims < 0:
    raise ValueError(
        'Input has fewer ndims ({}) than event shape ndims ({}).'.format(
            tensorshape_util.rank(input_tensorshape), event_shape_in_ndims))

  input_non_event_tensorshape = input_tensorshape[:input_non_event_ndims]
  input_event_tensorshape = input_tensorshape[input_non_event_ndims:]

  # Check that `input_event_shape_` and `event_shape_in` are compatible in the
  # sense that they have equal entries in any position that isn't a `-1` in
  # `event_shape_in`. Note that our validations at construction time ensure
  # there is at most one such entry in `event_shape_in`.
  event_shape_in_ = tf.get_static_value(event_shape_in)
  is_validated = (
      tensorshape_util.is_fully_defined(input_event_tensorshape) and
      event_shape_in_ is not None)
  if is_validated:
    input_event_shape_ = np.int32(input_event_tensorshape)
    mask = event_shape_in_ >= 0
    explicit_input_event_shape_ = input_event_shape_[mask]
    explicit_event_shape_in_ = event_shape_in_[mask]
    if not all(explicit_input_event_shape_ == explicit_event_shape_in_):
      raise ValueError(
          'Input `event_shape` does not match `event_shape_in`. '
          '({} vs {}).'.format(input_event_shape_, event_shape_in_))

  event_tensorshape_out = tensorshape_util.constant_value_as_shape(
      event_shape_out)
  if tensorshape_util.rank(event_tensorshape_out) is None:
    output_tensorshape = tf.TensorShape(None)
  else:
    output_tensorshape = tensorshape_util.concatenate(
        input_non_event_tensorshape, event_tensorshape_out)

  return output_tensorshape, is_validated

Replaces the event shape dims of a TensorShape .


def _maybe_check_valid_shape(shape, validate_args):
  """Check that a shape Tensor is int-type and otherwise sane."""
  if not dtype_util.is_integer(shape.dtype):
    raise TypeError('{} dtype ({}) should be `int`-like.'.format(
        shape, dtype_util.name(shape.dtype)))

  assertions = []

  message = '`{}` rank should be <= 1.'
  if tensorshape_util.rank(shape.shape) is not None:
    if tensorshape_util.rank(shape.shape) > 1:
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_less(
        tf.rank(shape), 2, message=message.format(shape)))

  shape_ = tf.get_static_value(shape)

  message = '`{}` elements must have at most one `-1`.'
  if shape_ is not None:
    if sum(shape_ == -1) > 1:
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_less(
        tf.reduce_sum(input_tensor=tf.cast(tf.equal(shape, -1), tf.int32)),
        2,
        message=message.format(shape)))

  message = '`{}` elements must be either positive integers or `-1`.'
  if shape_ is not None:
    if np.any(shape_ < -1):
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_greater(
        shape, -2, message=message.format(shape)))

  return assertions

Check that a shape Tensor is int - type and otherwise sane .


def converged_any(converged, failed):
  """Condition to stop when any batch member converges, or all have failed."""
  return (tf.reduce_any(input_tensor=converged) |
          tf.reduce_all(input_tensor=failed))

def converged_any ( converged , failed ) : return ( tf . reduce_any ( input_tensor = converged ) | tf . reduce_all ( input_tensor = failed ) )

Condition to stop when any batch member converges or all have failed .


def get_initial_state_args(value_and_gradients_function,
                           initial_position,
                           grad_tolerance,
                           control_inputs=None):
  """Returns a dictionary to populate the initial state of the search procedure.

  Performs an initial convergence check and the first evaluation of the
  objective function.

  Args:
    value_and_gradients_function: A Python callable that accepts a tensor and
      returns a tuple of two tensors: the objective function value and its
      derivative.
    initial_position: The starting point of the search procedure.
    grad_tolerance: The gradient tolerance for the procedure.
    control_inputs: Optional ops used to assert the validity of inputs, these
      are added as control dependencies to execute before the objective
      function is evaluated for the first time.

  Returns:
    An dictionary with values for the following keys:
      converged: True if the convergence check finds that the initial position
        is already an argmin of the objective function.
      failed: Initialized to False.
      num_objective_evaluations: Initialized to 1.
      position: Initialized to the initial position.
      objective_value: Initialized to the value of the objective function at
        the initial position.
      objective_gradient: Initialized to the gradient of the objective
        function at the initial position.
  """
  if control_inputs:
    with tf.control_dependencies(control_inputs):
      f0, df0 = value_and_gradients_function(initial_position)
  else:
    f0, df0 = value_and_gradients_function(initial_position)
  converged = norm(df0, dims=1) < grad_tolerance
  return dict(
      converged=converged,
      failed=tf.zeros_like(converged),  # i.e. False.
      num_iterations=tf.convert_to_tensor(value=0),
      num_objective_evaluations=tf.convert_to_tensor(value=1),
      position=initial_position,
      objective_value=f0,
      objective_gradient=df0)

def get_initial_state_args ( value_and_gradients_function , initial_position , grad_tolerance , control_inputs = None ) : if control_inputs : with tf . control_dependencies ( control_inputs ) : f0 , df0 = value_and_gradients_function ( initial_position ) else : f0 , df0 = value_and_gradients_function ( initial_position ) converged = norm ( df0 , dims = 1 ) < grad_tolerance return dict ( converged = converged , failed = tf . zeros_like ( converged ) , num_iterations = tf . convert_to_tensor ( value = 0 ) , num_objective_evaluations = tf . convert_to_tensor ( value = 1 ) , position = initial_position , objective_value = f0 , objective_gradient = df0 )

Returns a dictionary to populate the initial state of the search procedure .