tensorflow_probability/python/sts/internal/util.py
def _maybe_expand_trailing_dim(observed_time_series_tensor):
  """Ensures `observed_time_series_tensor` has a trailing dimension of size 1.

  The `tfd.LinearGaussianStateSpaceModel` Distribution has event shape of
  `[num_timesteps, observation_size]`, but canonical BSTS models
  are univariate, so their observation_size is always `1`. The extra trailing
  dimension gets annoying, so this method allows arguments with or without the
  extra dimension. There is no ambiguity except in the trivial special case
  where  `num_timesteps = 1`; this can be avoided by specifying any unit-length
  series in the explicit `[num_timesteps, 1]` style.

  Most users should not call this method directly, and instead call
  `canonicalize_observed_time_series_with_mask`, which handles converting
  to `Tensor` and specifying an optional missingness mask.

  Args:
    observed_time_series_tensor: `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` or `batch_shape + [num_timesteps]`,
      where `num_timesteps > 1`.

  Returns:
    expanded_time_series: `Tensor` of shape `batch_shape + [num_timesteps, 1]`.
  """

  with tf.compat.v1.name_scope(
      'maybe_expand_trailing_dim', values=[observed_time_series_tensor]):
    if (observed_time_series_tensor.shape.ndims is not None and
        tf.compat.dimension_value(
            observed_time_series_tensor.shape[-1]) is not None):
      expanded_time_series = (
          observed_time_series_tensor
          if observed_time_series_tensor.shape[-1] == 1
          else observed_time_series_tensor[..., tf.newaxis])
    else:
      expanded_time_series = tf.cond(
          pred=tf.equal(tf.shape(input=observed_time_series_tensor)[-1], 1),
          true_fn=lambda: observed_time_series_tensor,
          false_fn=lambda: observed_time_series_tensor[..., tf.newaxis])
    return expanded_time_series
def _maybe_expand_trailing_dim(observed_time_series_tensor):
  with tf.compat.v1.name_scope(
      'maybe_expand_trailing_dim', values=[observed_time_series_tensor]):
    if (observed_time_series_tensor.shape.ndims is not None and
        tf.compat.dimension_value(
            observed_time_series_tensor.shape[-1]) is not None):
      expanded_time_series = (
          observed_time_series_tensor
          if observed_time_series_tensor.shape[-1] == 1
          else observed_time_series_tensor[..., tf.newaxis])
    else:
      expanded_time_series = tf.cond(
          pred=tf.equal(tf.shape(input=observed_time_series_tensor)[-1], 1),
          true_fn=lambda: observed_time_series_tensor,
          false_fn=lambda: observed_time_series_tensor[..., tf.newaxis])
    return expanded_time_series
def _maybe_expand_trailing_dim ( observed_time_series_tensor ) : with tf . compat . v1 . name_scope ( 'maybe_expand_trailing_dim' , values = [ observed_time_series_tensor ] ) : if ( observed_time_series_tensor . shape . ndims is not None and tf . compat . dimension_value ( observed_time_series_tensor . shape [ - 1 ] ) is not None ) : expanded_time_series = ( observed_time_series_tensor if observed_time_series_tensor . shape [ - 1 ] == 1 else observed_time_series_tensor [ ... , tf . newaxis ] ) else : expanded_time_series = tf . cond ( pred = tf . equal ( tf . shape ( input = observed_time_series_tensor ) [ - 1 ] , 1 ) , true_fn = lambda : observed_time_series_tensor , false_fn = lambda : observed_time_series_tensor [ ... , tf . newaxis ] ) return expanded_time_series
Ensures `observed_time_series_tensor` has a trailing dimension of size 1.

  The `tfd.LinearGaussianStateSpaceModel` Distribution has event shape of
  `[num_timesteps, observation_size]`, but canonical BSTS models
  are univariate, so their observation_size is always `1`. The extra trailing
  dimension gets annoying, so this method allows arguments with or without the
  extra dimension. There is no ambiguity except in the trivial special case
  where  `num_timesteps = 1`; this can be avoided by specifying any unit-length
  series in the explicit `[num_timesteps, 1]` style.

  Most users should not call this method directly, and instead call
  `canonicalize_observed_time_series_with_mask`, which handles converting
  to `Tensor` and specifying an optional missingness mask.

  Args:
    observed_time_series_tensor: `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` or `batch_shape + [num_timesteps]`,
      where `num_timesteps > 1`.

  Returns:
    expanded_time_series: `Tensor` of shape `batch_shape + [num_timesteps, 1]`.
Ensures observed_time_series_tensor has a trailing dimension of size 1 .