def mixture_stddev(mixture_weight_vector, mean_vector, stddev_vector):
  """Computes the standard deviation of a mixture distribution.

  This function works regardless of the component distribution, so long as
  each component's mean and standard deviation can be provided.

  Args:
    mixture_weight_vector: A 2D tensor with shape [batch_size, num_components]
    mean_vector: A 2D tensor of mixture component means. Has shape `[batch_size,
      num_components]`.
    stddev_vector: A 2D tensor of mixture component standard deviations. Has
      shape `[batch_size, num_components]`.

  Returns:
    A 1D tensor of shape `[batch_size]` representing the standard deviation of
    the mixture distribution with given weights and component means and standard
    deviations.
  Raises:
    ValueError: If the shapes of the input tensors are not as expected.
  """
  tensorshape_util.assert_has_rank(mixture_weight_vector.shape, 2)
  if not tensorshape_util.is_compatible_with(mean_vector.shape,
                                             mixture_weight_vector.shape):
    raise ValueError("Expecting means to have same shape as mixture weights.")
  if not tensorshape_util.is_compatible_with(stddev_vector.shape,
                                             mixture_weight_vector.shape):
    raise ValueError("Expecting stddevs to have same shape as mixture weights.")

  # Reshape the distribution parameters for batched vectorized dot products.
  pi_for_dot_prod = tf.expand_dims(mixture_weight_vector, axis=1)
  mu_for_dot_prod = tf.expand_dims(mean_vector, axis=2)
  sigma_for_dot_prod = tf.expand_dims(stddev_vector, axis=2)

  # weighted average of component means under mixture distribution.
  mean_wa = tf.matmul(pi_for_dot_prod, mu_for_dot_prod)
  mean_wa = tf.reshape(mean_wa, (-1,))
  # weighted average of component variances under mixture distribution.
  var_wa = tf.matmul(pi_for_dot_prod, tf.square(sigma_for_dot_prod))
  var_wa = tf.reshape(var_wa, (-1,))
  # weighted average of component squared means under mixture distribution.
  sq_mean_wa = tf.matmul(pi_for_dot_prod, tf.square(mu_for_dot_prod))
  sq_mean_wa = tf.reshape(sq_mean_wa, (-1,))
  mixture_variance = var_wa + sq_mean_wa - tf.square(mean_wa)
  return tf.sqrt(mixture_variance)

Computes the standard deviation of a mixture distribution .


def make_tril_scale(loc=None,
                    scale_tril=None,
                    scale_diag=None,
                    scale_identity_multiplier=None,
                    shape_hint=None,
                    validate_args=False,
                    assert_positive=False,
                    name=None):
  """Creates a LinearOperator representing a lower triangular matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_tril: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k, k], which represents a k x k lower
      triangular matrix. When `None` no `scale_tril` term is added to the
      LinearOperator. The upper triangular elements above the diagonal are
      ignored.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
  """

  def _maybe_attach_assertion(x):
    if not validate_args:
      return x
    if assert_positive:
      return with_dependencies([
          assert_util.assert_positive(
              tf.linalg.diag_part(x), message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            tf.linalg.diag_part(x),
            tf.zeros([], x.dtype),
            message="diagonal part must be non-zero"),
    ], x)

  with tf.name_scope(name or "make_tril_scale"):

    dtype = dtype_util.common_dtype(
        [loc, scale_tril, scale_diag, scale_identity_multiplier],
        preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", dtype=dtype)
    scale_tril = _convert_to_tensor(scale_tril, name="scale_tril", dtype=dtype)
    scale_diag = _convert_to_tensor(scale_diag, name="scale_diag", dtype=dtype)
    scale_identity_multiplier = _convert_to_tensor(
        scale_identity_multiplier,
        name="scale_identity_multiplier",
        dtype=dtype)

  if scale_tril is not None:
    scale_tril = tf.linalg.band_part(scale_tril, -1, 0)  # Zero out TriU.
    tril_diag = tf.linalg.diag_part(scale_tril)
    if scale_diag is not None:
      tril_diag += scale_diag
    if scale_identity_multiplier is not None:
      tril_diag += scale_identity_multiplier[..., tf.newaxis]

    scale_tril = tf.linalg.set_diag(scale_tril, tril_diag)

    return tf.linalg.LinearOperatorLowerTriangular(
        tril=_maybe_attach_assertion(scale_tril),
        is_non_singular=True,
        is_self_adjoint=False,
        is_positive_definite=assert_positive)

  return make_diag_scale(
      loc=loc,
      scale_diag=scale_diag,
      scale_identity_multiplier=scale_identity_multiplier,
      shape_hint=shape_hint,
      validate_args=validate_args,
      assert_positive=assert_positive,
      name=name)

def
make_tril_scale
(
loc
=
,
scale_tril
=
,
scale_diag
=
,
scale_identity_multiplier
=
,
shape_hint
=
,
validate_args
=
False
,
assert_positive
=
False
,
name
=
)
:
def
_maybe_attach_assertion
(
x
)
:
if
not
validate_args
:
return
x
if
assert_positive
:
return
with_dependencies
(
[
assert_util
.
assert_positive
(
tf
.
linalg
.
diag_part
(
x
)
,
message
=
"diagonal part must be positive"
)
,
]
,
x
)
return
with_dependencies
(
[
assert_util
.
assert_none_equal
(
tf
.
linalg
.
diag_part
(
x
)
,
tf
.
zeros
(
[
]
,
x
.
dtype
)
,
message
=
"diagonal part must be non-zero"
)
,
]
,
x
)
with
tf
.
name_scope
(
name
or
"make_tril_scale"
)
:
dtype
=
dtype_util
.
common_dtype
(
[
loc
,
scale_tril
,
scale_diag
,
scale_identity_multiplier
]
,
preferred_dtype
=
tf
.
float32
)
loc
=
_convert_to_tensor
(
loc
,
name
=
"loc"
,
dtype
=
dtype
)
scale_tril
=
_convert_to_tensor
(
scale_tril
,
name
=
"scale_tril"
,
dtype
=
dtype
)
scale_diag
=
_convert_to_tensor
(
scale_diag
,
name
=
"scale_diag"
,
dtype
=
dtype
)
scale_identity_multiplier
=
_convert_to_tensor
(
scale_identity_multiplier
,
name
=
"scale_identity_multiplier"
,
dtype
=
dtype
)
if
scale_tril
is
not
:
scale_tril
=
tf
.
linalg
.
band_part
(
scale_tril
,
1
,
0
)
# Zero out TriU.
tril_diag
=
tf
.
linalg
.
diag_part
(
scale_tril
)
if
scale_diag
is
not
:
tril_diag
+=
scale_diag
if
scale_identity_multiplier
is
not
:
tril_diag
+=
scale_identity_multiplier
[
...
,
tf
.
newaxis
]
scale_tril
=
tf
.
linalg
.
set_diag
(
scale_tril
,
tril_diag
)
return
tf
.
linalg
.
LinearOperatorLowerTriangular
(
tril
=
_maybe_attach_assertion
(
scale_tril
)
,
is_non_singular
=
True
,
is_self_adjoint
=
False
,
is_positive_definite
=
assert_positive
)
return
make_diag_scale
(
loc
=
loc
,
scale_diag
=
scale_diag
,
scale_identity_multiplier
=
scale_identity_multiplier
,
shape_hint
=
shape_hint
,
validate_args
=
validate_args
,
assert_positive
=
assert_positive
,
name
=
name
)

Creates a LinearOperator representing a lower triangular matrix .


def make_diag_scale(loc=None,
                    scale_diag=None,
                    scale_identity_multiplier=None,
                    shape_hint=None,
                    validate_args=False,
                    assert_positive=False,
                    name=None,
                    dtype=None):
  """Creates a LinearOperator representing a diagonal matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.
    dtype: TF `DType` to prefer when converting args to `Tensor`s. Else, we fall
      back to a compatible dtype across all of `loc`, `scale_diag`, and
      `scale_identity_multiplier`.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
  """

  def _maybe_attach_assertion(x):
    if not validate_args:
      return x
    if assert_positive:
      return with_dependencies([
          assert_util.assert_positive(
              x, message="diagonal part must be positive"),
      ], x)
    return with_dependencies([
        assert_util.assert_none_equal(
            x, tf.zeros([], x.dtype), message="diagonal part must be non-zero")
    ], x)

  with tf.name_scope(name or "make_diag_scale"):
    if dtype is None:
      dtype = dtype_util.common_dtype(
          [loc, scale_diag, scale_identity_multiplier],
          preferred_dtype=tf.float32)
    loc = _convert_to_tensor(loc, name="loc", dtype=dtype)
    scale_diag = _convert_to_tensor(scale_diag, name="scale_diag", dtype=dtype)
    scale_identity_multiplier = _convert_to_tensor(
        scale_identity_multiplier,
        name="scale_identity_multiplier",
        dtype=dtype)

    if scale_diag is not None:
      if scale_identity_multiplier is not None:
        scale_diag += scale_identity_multiplier[..., tf.newaxis]
      return tf.linalg.LinearOperatorDiag(
          diag=_maybe_attach_assertion(scale_diag),
          is_non_singular=True,
          is_self_adjoint=True,
          is_positive_definite=assert_positive)

    if loc is None and shape_hint is None:
      raise ValueError("Cannot infer `event_shape` unless `loc` or "
                       "`shape_hint` is specified.")

    num_rows = shape_hint
    del shape_hint
    if num_rows is None:
      num_rows = tf.compat.dimension_value(loc.shape[-1])
      if num_rows is None:
        num_rows = tf.shape(input=loc)[-1]

    if scale_identity_multiplier is None:
      return tf.linalg.LinearOperatorIdentity(
          num_rows=num_rows,
          dtype=dtype,
          is_self_adjoint=True,
          is_positive_definite=True,
          assert_proper_shapes=validate_args)

    return tf.linalg.LinearOperatorScaledIdentity(
        num_rows=num_rows,
        multiplier=_maybe_attach_assertion(scale_identity_multiplier),
        is_non_singular=True,
        is_self_adjoint=True,
        is_positive_definite=assert_positive,
        assert_proper_shapes=validate_args)

Creates a LinearOperator representing a diagonal matrix .


def shapes_from_loc_and_scale(loc, scale, name="shapes_from_loc_and_scale"):
  """Infer distribution batch and event shapes from a location and scale.

  Location and scale family distributions determine their batch/event shape by
  broadcasting the `loc` and `scale` args.  This helper does that broadcast,
  statically if possible.

  Batch shape broadcasts as per the normal rules.
  We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
  allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
  must either be size `1`, or the same as `scale.range_dimension`.

  See `MultivariateNormalLinearOperator` for a usage example.

  Args:
    loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
      `rank(loc)==0`, both batch and event shape are determined by `scale`.
    scale:  A `LinearOperator` instance.
    name:  A string name to prepend to created ops.

  Returns:
    batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
    event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

  Raises:
    ValueError:  If the last dimension of `loc` is determined statically to be
      different than the range of `scale`.
  """
  if loc is not None and tensorshape_util.rank(loc.shape) == 0:
    loc = None  # scalar loc is irrelevant to determining batch/event shape.
  with tf.name_scope(name):
    # Get event shape.
    event_size = scale.range_dimension_tensor()
    event_size_ = tf.get_static_value(event_size)
    loc_event_size_ = (None if loc is None
                       else tf.compat.dimension_value(loc.shape[-1]))

    if event_size_ is not None and loc_event_size_ is not None:
      # Static check that event shapes match.
      if loc_event_size_ != 1 and loc_event_size_ != event_size_:
        raise ValueError(
            "Event size of 'scale' ({}) could not be broadcast up to that "
            "of 'loc' ({}).".format(event_size_, loc_event_size_))
    elif loc_event_size_ is not None and loc_event_size_ != 1:
      event_size_ = loc_event_size_

    if event_size_ is None:
      event_shape = event_size[tf.newaxis]
    else:
      event_shape = tf.convert_to_tensor(
          value=np.reshape(event_size_, [1]),
          dtype=tf.int32,
          name="event_shape")

    # Get batch shape.
    batch_shape = scale.batch_shape_tensor()
    if loc is not None:
      loc_batch_shape = tensorshape_util.with_rank_at_least(loc.shape, 1)[:-1]
      if tensorshape_util.rank(
          loc.shape) is None or not tensorshape_util.is_fully_defined(
              loc_batch_shape):
        loc_batch_shape = tf.shape(input=loc)[:-1]
      else:
        loc_batch_shape = tf.convert_to_tensor(
            value=loc_batch_shape, dtype=tf.int32, name="loc_batch_shape")
      # This is defined in the core util module.
      batch_shape = prefer_static_broadcast_shape(batch_shape, loc_batch_shape)  # pylint: disable=undefined-variable
      batch_shape = tf.convert_to_tensor(
          value=batch_shape, dtype=tf.int32, name="batch_shape")

    return batch_shape, event_shape

Infer distribution batch and event shapes from a location and scale .


def is_diagonal_scale(scale):
  """Returns `True` if `scale` is a `LinearOperator` that is known to be diag.

  Args:
    scale:  `LinearOperator` instance.

  Returns:
    Python `bool`.

  Raises:
    TypeError:  If `scale` is not a `LinearOperator`.
  """
  if not isinstance(scale, tf.linalg.LinearOperator):
    raise TypeError("Expected argument 'scale' to be instance of LinearOperator"
                    ". Found: %s" % scale)
  return (isinstance(scale, tf.linalg.LinearOperatorIdentity) or
          isinstance(scale, tf.linalg.LinearOperatorScaledIdentity) or
          isinstance(scale, tf.linalg.LinearOperatorDiag))

Returns True if scale is a LinearOperator that is known to be diag .


def maybe_check_scalar_distribution(distribution, expected_base_dtype,
                                    validate_args):
  """Helper which checks validity of a scalar `distribution` init arg.

  Valid here means:

  * `distribution` has scalar batch and event shapes.
  * `distribution` is `FULLY_REPARAMETERIZED`
  * `distribution` has expected dtype.

  Args:
    distribution:  `Distribution`-like object.
    expected_base_dtype:  `TensorFlow` `dtype`.
    validate_args:  Python `bool`.  Whether to do additional checks: (i)  check
      that reparameterization_type is `FULLY_REPARAMETERIZED`. (ii) add
      `tf.Assert` ops to the graph to enforce that distribution is scalar in the
      event that this cannot be determined statically.

  Returns:
    List of `tf.Assert` ops to run to enforce validity checks that could not
      be statically determined.  Empty if `not validate_args`.

  Raises:
    ValueError:  If validate_args and distribution is not FULLY_REPARAMETERIZED
    ValueError:  If distribution is statically determined to not have both
      scalar batch and scalar event shapes.
  """
  if distribution.dtype != expected_base_dtype:
    raise TypeError("dtype mismatch; "
                    "distribution.dtype=\"{}\" is not \"{}\"".format(
                        dtype_util.name(distribution.dtype),
                        dtype_util.name(expected_base_dtype)))

  # Although `reparameterization_type` is a static property, we guard it by
  # `validate_args`. This allows users to use a `distribution` which is not
  # reparameterized itself. However, we tacitly assume that although the
  # distribution is not reparameterized, it only depends on non-trainable
  # variables.
  if validate_args and (distribution.reparameterization_type !=
                        reparameterization.FULLY_REPARAMETERIZED):
    raise ValueError("Base distribution should be reparameterized or be "
                     "a function of non-trainable variables; "
                     "distribution.reparameterization_type = \"{}\" "
                     "!= \"FULLY_REPARAMETERIZED\".".format(
                         distribution.reparameterization_type))
  with tf.name_scope("check_distribution"):
    assertions = []

    def check_is_scalar(is_scalar, name):
      is_scalar_ = tf.get_static_value(is_scalar)
      if is_scalar_ is not None:
        if not is_scalar_:
          raise ValueError("distribution must be scalar; "
                           "distribution.{}=False is not True".format(name))
      elif validate_args:
        assertions.append(
            assert_util.assert_equal(
                is_scalar,
                True,
                message=("distribution must be scalar; "
                         "distribution.{}=False is not True".format(name))))

    check_is_scalar(distribution.is_scalar_event(), "is_scalar_event")
    check_is_scalar(distribution.is_scalar_batch(), "is_scalar_batch")
    return assertions

Helper which checks validity of a scalar distribution init arg .


def pad_mixture_dimensions(x, mixture_distribution, categorical_distribution,
                           event_ndims):
  """Pad dimensions of event tensors for mixture distributions.

  See `Mixture._sample_n` and `MixtureSameFamily._sample_n` for usage examples.

  Args:
    x: event tensor to pad.
    mixture_distribution: Base distribution of the mixture.
    categorical_distribution: `Categorical` distribution that mixes the base
      distribution.
    event_ndims: Integer specifying the number of event dimensions in the event
      tensor.

  Returns:
    A padded version of `x` that can broadcast with `categorical_distribution`.
  """
  with tf.name_scope("pad_mix_dims"):

    def _get_ndims(d):
      if tensorshape_util.rank(d.batch_shape) is not None:
        return tensorshape_util.rank(d.batch_shape)
      return tf.shape(input=d.batch_shape_tensor())[0]

    dist_batch_ndims = _get_ndims(mixture_distribution)
    cat_batch_ndims = _get_ndims(categorical_distribution)
    pad_ndims = tf.where(categorical_distribution.is_scalar_batch(),
                         dist_batch_ndims, dist_batch_ndims - cat_batch_ndims)
    s = tf.shape(input=x)
    x = tf.reshape(
        x,
        shape=tf.concat([
            s[:-1],
            tf.ones([pad_ndims], dtype=tf.int32),
            s[-1:],
            tf.ones([event_ndims], dtype=tf.int32),
        ],
                        axis=0))
    return x

def pad_mixture_dimensions ( x , mixture_distribution , categorical_distribution , event_ndims ) : with tf . name_scope ( "pad_mix_dims" ) : def _get_ndims ( d ) : if tensorshape_util . rank ( d . batch_shape ) is not None : return tensorshape_util . rank ( d . batch_shape ) return tf . shape ( input = d . batch_shape_tensor ( ) ) [ 0 ] dist_batch_ndims = _get_ndims ( mixture_distribution ) cat_batch_ndims = _get_ndims ( categorical_distribution ) pad_ndims = tf . where ( categorical_distribution . is_scalar_batch ( ) , dist_batch_ndims , dist_batch_ndims - cat_batch_ndims ) s = tf . shape ( input = x ) x = tf . reshape ( x , shape = tf . concat ( [ s [ : - 1 ] , tf . ones ( [ pad_ndims ] , dtype = tf . int32 ) , s [ - 1 : ] , tf . ones ( [ event_ndims ] , dtype = tf . int32 ) , ] , axis = 0 ) ) return x

Pad dimensions of event tensors for mixture distributions .


def pick_scalar_condition(pred, true_value, false_value, name=None):
  """Convenience function that chooses one of two values based on the predicate.

  This utility is equivalent to a version of `tf.where` that accepts only a
  scalar predicate and computes its result statically when possible. It may also
  be used in place of `tf.cond` when both branches yield a `Tensor` of the same
  shape; the operational difference is that `tf.cond` uses control flow to
  evaluate only the branch that's needed, while `tf.where` (and thus
  this method) may evaluate both branches before the predicate's truth is known.
  This means that `tf.cond` is preferred when one of the branches is expensive
  to evaluate (like performing a large matmul), while this method is preferred
  when both branches are cheap, e.g., constants. In the latter case, we expect
  this method to be substantially faster than `tf.cond` on GPU and to give
  similar performance on CPU.

  Args:
    pred: Scalar `bool` `Tensor` predicate.
    true_value: `Tensor` to return if `pred` is `True`.
    false_value: `Tensor` to return if `pred` is `False`. Must have the same
      shape as `true_value`.
    name: Python `str` name given to ops managed by this object.

  Returns:
    result: a `Tensor` (or `Tensor`-convertible Python value) equal to
      `true_value` if `pred` evaluates to `True` and `false_value` otherwise.
      If the condition can be evaluated statically, the result returned is one
      of the input Python values, with no graph side effects.
  """
  with tf.name_scope(name or "pick_scalar_condition"):
    pred = tf.convert_to_tensor(
        value=pred, dtype_hint=tf.bool, name="pred")
    true_value = tf.convert_to_tensor(value=true_value, name="true_value")
    false_value = tf.convert_to_tensor(value=false_value, name="false_value")
    pred_ = tf.get_static_value(pred)
    if pred_ is None:
      return tf.where(pred, true_value, false_value)
    return true_value if pred_ else false_value

def pick_scalar_condition ( pred , true_value , false_value , name = None ) : with tf . name_scope ( name or "pick_scalar_condition" ) : pred = tf . convert_to_tensor ( value = pred , dtype_hint = tf . bool , name = "pred" ) true_value = tf . convert_to_tensor ( value = true_value , name = "true_value" ) false_value = tf . convert_to_tensor ( value = false_value , name = "false_value" ) pred_ = tf . get_static_value ( pred ) if pred_ is None : return tf . where ( pred , true_value , false_value ) return true_value if pred_ else false_value

Convenience function that chooses one of two values based on the predicate .


def move_dimension(x, source_idx, dest_idx):
  """Move a single tensor dimension within its shape.

  This is a special case of `tf.transpose()`, which applies
  arbitrary permutations to tensor dimensions.

  Args:
    x: Tensor of rank `ndims`.
    source_idx: Integer index into `x.shape` (negative indexing is supported).
    dest_idx: Integer index into `x.shape` (negative indexing is supported).

  Returns:
    x_perm: Tensor of rank `ndims`, in which the dimension at original
     index `source_idx` has been moved to new index `dest_idx`, with
     all other dimensions retained in their original order.

  Example:

  ```python
  x = tf.placeholder(shape=[200, 30, 4, 1, 6])
  x_perm = _move_dimension(x, 1, 1) # no-op
  x_perm = _move_dimension(x, 0, 3) # result shape [30, 4, 1, 200, 6]
  x_perm = _move_dimension(x, 0, -2) # equivalent to previous
  x_perm = _move_dimension(x, 4, 2) # result shape [200, 30, 6, 4, 1]
  ```
  """
  ndims = prefer_static_rank(x)
  dtype = dtype_util.common_dtype([source_idx, dest_idx],
                                  preferred_dtype=tf.int32)
  source_idx = tf.convert_to_tensor(value=source_idx, dtype=dtype)
  dest_idx = tf.convert_to_tensor(value=dest_idx, dtype=dtype)

  # Handle negative indexing.
  source_idx = pick_scalar_condition(source_idx < 0, ndims + source_idx,
                                     source_idx)
  dest_idx = pick_scalar_condition(dest_idx < 0, ndims + dest_idx, dest_idx)

  # Construct the appropriate permutation of dimensions, depending
  # whether the source is before or after the destination.
  def move_left_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, dest_idx, dtype=dtype), [source_idx],
            tf.range(dest_idx, source_idx, dtype=dtype),
            tf.range(source_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def move_right_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, source_idx, dtype=dtype),
            tf.range(source_idx + 1, dest_idx + 1, dtype=dtype), [source_idx],
            tf.range(dest_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def x_permuted():
    return tf.transpose(
        a=x,
        perm=prefer_static.cond(source_idx < dest_idx,
                                move_right_permutation,
                                move_left_permutation))

  # One final conditional to handle the special case where source
  # and destination indices are equal.
  return prefer_static.cond(tf.equal(source_idx, dest_idx),
                            lambda: x, x_permuted)

Move a single tensor dimension within its shape .


def embed_check_nonnegative_integer_form(
    x, name="embed_check_nonnegative_integer_form"):
  """Assert x is a non-negative tensor, and optionally of integers."""
  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    assertions = [
        assert_util.assert_non_negative(
            x, message="'{}' must be non-negative.".format(x)),
    ]
    if not dtype_util.is_integer(x.dtype):
      assertions += [
          assert_integer_form(
              x,
              message="'{}' cannot contain fractional components.".format(x)),
      ]
    return with_dependencies(assertions, x)

Assert x is a non - negative tensor and optionally of integers .


def same_dynamic_shape(a, b):
  """Returns whether a and b have the same dynamic shape.

  Args:
    a: `Tensor`
    b: `Tensor`

  Returns:
    `bool` `Tensor` representing if both tensors have the same shape.
  """
  a = tf.convert_to_tensor(value=a, name="a")
  b = tf.convert_to_tensor(value=b, name="b")

  # Here we can't just do tf.equal(a.shape, b.shape), since
  # static shape inference may break the equality comparison between
  # shape(a) and shape(b) in tf.equal.
  def all_shapes_equal():
    return tf.reduce_all(
        input_tensor=tf.equal(
            tf.concat([tf.shape(input=a), tf.shape(input=b)], 0),
            tf.concat([tf.shape(input=b), tf.shape(input=a)], 0)))

  # One of the shapes isn't fully defined, so we need to use the dynamic
  # shape.
  return tf.cond(
      pred=tf.equal(tf.rank(a), tf.rank(b)),
      true_fn=all_shapes_equal,
      false_fn=lambda: tf.constant(False))

Returns whether a and b have the same dynamic shape .


def maybe_get_static_value(x, dtype=None):
  """Helper which tries to return a static value.

  Given `x`, extract it's value statically, optionally casting to a specific
  dtype. If this is not possible, None is returned.

  Args:
    x: `Tensor` for which to extract a value statically.
    dtype: Optional dtype to cast to.

  Returns:
    Statically inferred value if possible, otherwise None.
  """
  if x is None:
    return x
  try:
    # This returns an np.ndarray.
    x_ = tf.get_static_value(x)
  except TypeError:
    x_ = x
  if x_ is None or dtype is None:
    return x_
  return np.array(x_, dtype)

def
maybe_get_static_value
(
x
,
dtype
=
)
:
if
x
is
:
return
x
try
:
# This returns an np.ndarray.
x_
=
tf
.
get_static_value
(
x
)
except
TypeError
:
x_
=
x
if
x_
is
or
dtype
is
:
return
x_
return
np
.
array
(
x_
,
dtype
)

Helper which tries to return a static value .


def _is_known_unsigned_by_dtype(dt):
  """Helper returning True if dtype is known to be unsigned."""
  return {
      tf.bool: True,
      tf.uint8: True,
      tf.uint16: True,
  }.get(dt.base_dtype, False)

def _is_known_unsigned_by_dtype ( dt ) : return { tf . bool : True , tf . uint8 : True , tf . uint16 : True , } . get ( dt . base_dtype , False )

Helper returning True if dtype is known to be unsigned .


def _is_known_signed_by_dtype(dt):
  """Helper returning True if dtype is known to be signed."""
  return {
      tf.float16: True,
      tf.float32: True,
      tf.float64: True,
      tf.int8: True,
      tf.int16: True,
      tf.int32: True,
      tf.int64: True,
  }.get(dt.base_dtype, False)

def _is_known_signed_by_dtype ( dt ) : return { tf . float16 : True , tf . float32 : True , tf . float64 : True , tf . int8 : True , tf . int16 : True , tf . int32 : True , tf . int64 : True , } . get ( dt . base_dtype , False )

Helper returning True if dtype is known to be signed .


def _largest_integer_by_dtype(dt):
  """Helper returning the largest integer exactly representable by dtype."""
  if not _is_known_dtype(dt):
    raise TypeError("Unrecognized dtype: {}".format(dt.name))
  if dt.is_floating:
    return int(2**(np.finfo(dt.as_numpy_dtype).nmant + 1))
  if dt.is_integer:
    return np.iinfo(dt.as_numpy_dtype).max
  if dt.base_dtype == tf.bool:
    return int(1)
  # We actually can't land here but keep the case for completeness.
  raise TypeError("Unrecognized dtype: {}".format(dt.name))

Helper returning the largest integer exactly representable by dtype .


def _smallest_integer_by_dtype(dt):
  """Helper returning the smallest integer exactly representable by dtype."""
  if not _is_known_dtype(dt):
    raise TypeError("Unrecognized dtype: {}".format(dt.name))
  if _is_known_unsigned_by_dtype(dt):
    return 0
  return -1 * _largest_integer_by_dtype(dt)

def
_smallest_integer_by_dtype
(
dt
)
:
if
not
_is_known_dtype
(
dt
)
:
raise
TypeError
(
"Unrecognized dtype: {}"
.
format
(
dt
.
name
)
)
if
_is_known_unsigned_by_dtype
(
dt
)
:
return
0
return
1
*
_largest_integer_by_dtype
(
dt
)

Helper returning the smallest integer exactly representable by dtype .


def _is_integer_like_by_dtype(dt):
  """Helper returning True if dtype.is_integer or is `bool`."""
  if not _is_known_dtype(dt):
    raise TypeError("Unrecognized dtype: {}".format(dt.name))
  return dt.is_integer or dt.base_dtype == tf.bool

def
_is_integer_like_by_dtype
(
dt
)
:
if
not
_is_known_dtype
(
dt
)
:
raise
TypeError
(
"Unrecognized dtype: {}"
.
format
(
dt
.
name
)
)
return
dt
.
is_integer
or
dt
.
base_dtype
==
tf
.
bool

Helper returning True if dtype . is_integer or is bool .


def embed_check_categorical_event_shape(
    categorical_param, name="embed_check_categorical_event_shape"):
  """Embeds checks that categorical distributions don't have too many classes.

  A categorical-type distribution is one which, e.g., returns the class label
  rather than a one-hot encoding.  E.g., `Categorical(probs)`.

  Since distributions output samples in the same dtype as the parameters, we
  must ensure that casting doesn't lose precision. That is, the
  `parameter.dtype` implies a maximum number of classes. However, since shape is
  `int32` and categorical variables are presumed to be indexes into a `Tensor`,
  we must also ensure that the number of classes is no larger than the largest
  possible `int32` index, i.e., `2**31-1`.

  In other words the number of classes, `K`, must satisfy the following
  condition:

  ```python
  K <= min(
      int(2**31 - 1),  # Largest float as an index.
      {
          tf.float16: int(2**11),   # Largest int as a float16.
          tf.float32: int(2**24),
          tf.float64: int(2**53),
      }.get(dtype_util.base_dtype(categorical_param.dtype), 0))
  ```

  Args:
    categorical_param: Floating-point `Tensor` representing parameters of
      distribution over categories. The rightmost shape is presumed to be the
      number of categories.
    name: A name for this operation (optional).

  Returns:
    categorical_param: Input `Tensor` with appropriate assertions embedded.

  Raises:
    TypeError: if `categorical_param` has an unknown `dtype`.
    ValueError: if we can statically identify `categorical_param` as being too
      large (for being closed under int32/float casting).
  """
  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=categorical_param, name="categorical_param")
    # The size must not exceed both of:
    # - The largest possible int32 (since categorical values are presumed to be
    #   indexes into a Tensor).
    # - The largest possible integer exactly representable under the given
    #   floating-point dtype (since we need to cast to/from).
    #
    # The chosen floating-point thresholds are 2**(1 + mantissa_bits).
    # For more details, see:
    # https://en.wikipedia.org/wiki/Floating-point_arithmetic#Internal_representation
    x_dtype = dtype_util.base_dtype(x.dtype)
    max_event_size = (
        _largest_integer_by_dtype(x_dtype)
        if dtype_util.is_floating(x_dtype) else 0)
    if max_event_size is 0:
      raise TypeError("Unable to validate size of unrecognized dtype "
                      "({}).".format(dtype_util.name(x_dtype)))
    try:
      x_shape_static = tensorshape_util.with_rank_at_least(x.shape, 1)
    except ValueError:
      raise ValueError("A categorical-distribution parameter must have "
                       "at least 1 dimension.")
    event_size = tf.compat.dimension_value(x_shape_static[-1])
    if event_size is not None:
      if event_size < 2:
        raise ValueError("A categorical-distribution parameter must have at "
                         "least 2 events.")
      if event_size > max_event_size:
        raise ValueError("Number of classes exceeds `dtype` precision, i.e., "
                         "{} implies shape ({}) cannot exceed {}.".format(
                             dtype_util.name(x_dtype), event_size,
                             max_event_size))
      return x
    else:
      event_size = tf.shape(input=x, out_type=tf.int64, name="x_shape")[-1]
      return with_dependencies([
          assert_util.assert_rank_at_least(
              x,
              1,
              message=("A categorical-distribution parameter must have "
                       "at least 1 dimension.")),
          assert_util.assert_greater_equal(
              tf.shape(input=x)[-1],
              2,
              message=("A categorical-distribution parameter must have at "
                       "least 2 events.")),
          assert_util.assert_less_equal(
              event_size,
              tf.convert_to_tensor(max_event_size, dtype=tf.int64),
              message="Number of classes exceeds `dtype` precision, "
              "i.e., {} dtype cannot exceed {} shape.".format(
                  dtype_util.name(x_dtype), max_event_size)),
      ], x)

Embeds checks that categorical distributions don t have too many classes .


def log_combinations(n, counts, name="log_combinations"):
  """Multinomial coefficient.

  Given `n` and `counts`, where `counts` has last dimension `k`, we compute
  the multinomial coefficient as:

  ```n! / sum_i n_i!```

  where `i` runs over all `k` classes.

  Args:
    n: Floating-point `Tensor` broadcastable with `counts`. This represents `n`
      outcomes.
    counts: Floating-point `Tensor` broadcastable with `n`. This represents
      counts in `k` classes, where `k` is the last dimension of the tensor.
    name: A name for this operation (optional).

  Returns:
    `Tensor` representing the multinomial coefficient between `n` and `counts`.
  """
  # First a bit about the number of ways counts could have come in:
  # E.g. if counts = [1, 2], then this is 3 choose 2.
  # In general, this is (sum counts)! / sum(counts!)
  # The sum should be along the last dimension of counts. This is the
  # "distribution" dimension. Here n a priori represents the sum of counts.
  with tf.name_scope(name):
    n = tf.convert_to_tensor(value=n, name="n")
    counts = tf.convert_to_tensor(value=counts, name="counts")
    total_permutations = tf.math.lgamma(n + 1)
    counts_factorial = tf.math.lgamma(counts + 1)
    redundant_permutations = tf.reduce_sum(
        input_tensor=counts_factorial, axis=[-1])
    return total_permutations - redundant_permutations


def rotate_transpose(x, shift, name="rotate_transpose"):
  """Circularly moves dims left or right.

  Effectively identical to:

  ```python
  numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift))
  ```

  When `validate_args=False` additional graph-runtime checks are
  performed. These checks entail moving data from to GPU to CPU.

  Example:

  ```python
  x = tf.random_normal([1, 2, 3, 4])  # Tensor of shape [1, 2, 3, 4].
  rotate_transpose(x, -1).shape == [2, 3, 4, 1]
  rotate_transpose(x, -2).shape == [3, 4, 1, 2]
  rotate_transpose(x,  1).shape == [4, 1, 2, 3]
  rotate_transpose(x,  2).shape == [3, 4, 1, 2]
  rotate_transpose(x,  7).shape == rotate_transpose(x, 3).shape  # [2, 3, 4, 1]
  rotate_transpose(x, -7).shape == rotate_transpose(x, -3).shape  # [4, 1, 2, 3]
  ```

  Args:
    x: `Tensor`.
    shift: `Tensor`. Number of dimensions to transpose left (shift<0) or
      transpose right (shift>0).
    name: Python `str`. The name to give this op.

  Returns:
    rotated_x: Input `Tensor` with dimensions circularly rotated by shift.

  Raises:
    TypeError: if shift is not integer type.
  """
  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    shift = tf.convert_to_tensor(value=shift, name="shift")
    # We do not assign back to preserve constant-ness.
    assert_util.assert_integer(shift)
    shift_value_static = tf.get_static_value(shift)
    ndims = tensorshape_util.rank(x.shape)
    if ndims is not None and shift_value_static is not None:
      if ndims < 2:
        return x
      shift_value_static = np.sign(shift_value_static) * (
          abs(shift_value_static) % ndims)
      if shift_value_static == 0:
        return x
      perm = np.roll(np.arange(ndims), shift_value_static)
      return tf.transpose(a=x, perm=perm)
    else:
      # Consider if we always had a positive shift, and some specified
      # direction.
      # When shifting left we want the new array:
      #   last(x, n-shift) + first(x, shift)
      # and if shifting right then we want:
      #   last(x, shift) + first(x, n-shift)
      # Observe that last(a) == slice(a, n) and first(a) == slice(0, a).
      # Also, we can encode direction and shift as one: direction * shift.
      # Combining these facts, we have:
      #   a = cond(shift<0, -shift, n-shift)
      #   last(x, n-a) + first(x, a) == x[a:n] + x[0:a]
      # Finally, we transform shift by modulo length so it can be specified
      # independently from the array upon which it operates (like python).
      ndims = tf.rank(x)
      shift = tf.where(
          tf.less(shift, 0), -shift % ndims,
          ndims - shift % ndims)
      first = tf.range(0, shift)
      last = tf.range(shift, ndims)
      perm = tf.concat([last, first], 0)
      return tf.transpose(a=x, perm=perm)


def pick_vector(cond, true_vector, false_vector, name="pick_vector"):
  """Picks possibly different length row `Tensor`s based on condition.

  Value `Tensor`s should have exactly one dimension.

  If `cond` is a python Boolean or `tf.constant` then either `true_vector` or
  `false_vector` is immediately returned. I.e., no graph nodes are created and
  no validation happens.

  Args:
    cond: `Tensor`. Must have `dtype=tf.bool` and be scalar.
    true_vector: `Tensor` of one dimension. Returned when cond is `True`.
    false_vector: `Tensor` of one dimension. Returned when cond is `False`.
    name: Python `str`. The name to give this op.
  Example:  ```python pick_vector(tf.less(0, 5), tf.range(10, 12), tf.range(15,
    18))  # [10, 11] pick_vector(tf.less(5, 0), tf.range(10, 12), tf.range(15,
    18))  # [15, 16, 17] ```

  Returns:
    true_or_false_vector: `Tensor`.

  Raises:
    TypeError: if `cond.dtype != tf.bool`
    TypeError: if `cond` is not a constant and
      `true_vector.dtype != false_vector.dtype`
  """
  with tf.name_scope(name):
    cond = tf.convert_to_tensor(
        value=cond, dtype_hint=tf.bool, name="cond")
    if cond.dtype != tf.bool:
      raise TypeError(
          "{}.dtype={} which is not {}".format(cond, cond.dtype, tf.bool))

    true_vector = tf.convert_to_tensor(value=true_vector, name="true_vector")
    false_vector = tf.convert_to_tensor(value=false_vector, name="false_vector")
    if true_vector.dtype != false_vector.dtype:
      raise TypeError(
          "{}.dtype={} does not match {}.dtype={}".format(
              true_vector, true_vector.dtype, false_vector, false_vector.dtype))

    cond_value_static = tf.get_static_value(cond)
    if cond_value_static is not None:
      return true_vector if cond_value_static else false_vector
    n = tf.shape(input=true_vector)[0]
    return tf.slice(
        tf.concat([true_vector, false_vector], 0), [tf.where(cond, 0, n)],
        [tf.where(cond, n, -1)])

Picks possibly different length row Tensor s based on condition .


def prefer_static_broadcast_shape(shape1,
                                  shape2,
                                  name="prefer_static_broadcast_shape"):
  """Convenience function which statically broadcasts shape when possible.

  Args:
    shape1:  `1-D` integer `Tensor`.  Already converted to tensor!
    shape2:  `1-D` integer `Tensor`.  Already converted to tensor!
    name:  A string name to prepend to created ops.

  Returns:
    The broadcast shape, either as `TensorShape` (if broadcast can be done
      statically), or as a `Tensor`.
  """
  with tf.name_scope(name):

    def make_shape_tensor(x):
      return tf.convert_to_tensor(value=x, name="shape", dtype=tf.int32)

    def get_tensor_shape(s):
      if isinstance(s, tf.TensorShape):
        return s
      s_ = tf.get_static_value(make_shape_tensor(s))
      if s_ is not None:
        return tf.TensorShape(s_)
      return None

    def get_shape_tensor(s):
      if not isinstance(s, tf.TensorShape):
        return make_shape_tensor(s)
      if tensorshape_util.is_fully_defined(s):
        return make_shape_tensor(tensorshape_util.as_list(s))
      raise ValueError("Cannot broadcast from partially "
                       "defined `TensorShape`.")

    shape1_ = get_tensor_shape(shape1)
    shape2_ = get_tensor_shape(shape2)
    if shape1_ is not None and shape2_ is not None:
      return tf.broadcast_static_shape(shape1_, shape2_)

    shape1_ = get_shape_tensor(shape1)
    shape2_ = get_shape_tensor(shape2)
    return tf.broadcast_dynamic_shape(shape1_, shape2_)

def
prefer_static_broadcast_shape
(
shape1
,
shape2
,
name
=
"prefer_static_broadcast_shape"
)
:
with
tf
.
name_scope
(
name
)
:
def
make_shape_tensor
(
x
)
:
return
tf
.
convert_to_tensor
(
value
=
x
,
name
=
"shape"
,
dtype
=
tf
.
int32
)
def
get_tensor_shape
(
s
)
:
if
isinstance
(
s
,
tf
.
TensorShape
)
:
return
s
s_
=
tf
.
get_static_value
(
make_shape_tensor
(
s
)
)
if
s_
is
not
:
return
tf
.
TensorShape
(
s_
)
return
def
get_shape_tensor
(
s
)
:
if
not
isinstance
(
s
,
tf
.
TensorShape
)
:
return
make_shape_tensor
(
s
)
if
tensorshape_util
.
is_fully_defined
(
s
)
:
return
make_shape_tensor
(
tensorshape_util
.
as_list
(
s
)
)
raise
ValueError
(
"Cannot broadcast from partially "
"defined `TensorShape`."
)
shape1_
=
get_tensor_shape
(
shape1
)
shape2_
=
get_tensor_shape
(
shape2
)
if
shape1_
is
not
and
shape2_
is
not
:
return
tf
.
broadcast_static_shape
(
shape1_
,
shape2_
)
shape1_
=
get_shape_tensor
(
shape1
)
shape2_
=
get_shape_tensor
(
shape2
)
return
tf
.
broadcast_dynamic_shape
(
shape1_
,
shape2_
)

Convenience function which statically broadcasts shape when possible .


def gen_new_seed(seed, salt):
  """Generate a new seed, from the given seed and salt."""
  if seed is None:
    return None
  string = (str(seed) + salt).encode("utf-8")
  return int(hashlib.md5(string).hexdigest()[:8], 16) & 0x7FFFFFFF

def gen_new_seed ( seed , salt ) : if seed is None : return None string = ( str ( seed ) + salt ) . encode ( "utf-8" ) return int ( hashlib . md5 ( string ) . hexdigest ( ) [ : 8 ] , 16 ) & 0x7FFFFFFF

Generate a new seed from the given seed and salt .


def tridiag(below=None, diag=None, above=None, name=None):
  """Creates a matrix with values set above, below, and on the diagonal.

  Example:

  ```python
  tridiag(below=[1., 2., 3.],
          diag=[4., 5., 6., 7.],
          above=[8., 9., 10.])
  # ==> array([[  4.,   8.,   0.,   0.],
  #            [  1.,   5.,   9.,   0.],
  #            [  0.,   2.,   6.,  10.],
  #            [  0.,   0.,   3.,   7.]], dtype=float32)
  ```

  Warning: This Op is intended for convenience, not efficiency.

  Args:
    below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below
      diagonal part. `None` is logically equivalent to `below = 0`.
    diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal
      part.  `None` is logically equivalent to `diag = 0`.
    above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above
      diagonal part.  `None` is logically equivalent to `above = 0`.
    name: Python `str`. The name to give this op.

  Returns:
    tridiag: `Tensor` with values set above, below and on the diagonal.

  Raises:
    ValueError: if all inputs are `None`.
  """

  def _pad(x):
    """Prepends and appends a zero to every vector in a batch of vectors."""
    shape = tf.concat([tf.shape(input=x)[:-1], [1]], axis=0)
    z = tf.zeros(shape, dtype=x.dtype)
    return tf.concat([z, x, z], axis=-1)

  def _add(*x):
    """Adds list of Tensors, ignoring `None`."""
    s = None
    for y in x:
      if y is None:
        continue
      elif s is None:
        s = y
      else:
        s += y
    if s is None:
      raise ValueError("Must specify at least one of `below`, `diag`, `above`.")
    return s

  with tf.name_scope(name or "tridiag"):
    if below is not None:
      below = tf.convert_to_tensor(value=below, name="below")
      below = tf.linalg.diag(_pad(below))[..., :-1, 1:]
    if diag is not None:
      diag = tf.convert_to_tensor(value=diag, name="diag")
      diag = tf.linalg.diag(diag)
    if above is not None:
      above = tf.convert_to_tensor(value=above, name="above")
      above = tf.linalg.diag(_pad(above))[..., 1:, :-1]
    # TODO(jvdillon): Consider using scatter_nd instead of creating three full
    # matrices.
    return _add(below, diag, above)

Creates a matrix with values set above below and on the diagonal .


def dimension_size(x, axis):
  """Returns the size of a specific dimension."""
  # Since tf.gather isn't "constant-in, constant-out", we must first check the
  # static shape or fallback to dynamic shape.
  s = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(x.shape, np.abs(axis))[axis])
  if s is not None:
    return s
  return tf.shape(input=x)[axis]

Returns the size of a specific dimension .


def process_quadrature_grid_and_probs(quadrature_grid_and_probs,
                                      dtype,
                                      validate_args,
                                      name=None):
  """Validates quadrature grid, probs or computes them as necessary.

  Args:
    quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
      representing the sample points and the corresponding (possibly
      normalized) weight.  When `None`, defaults to:
        `np.polynomial.hermite.hermgauss(deg=8)`.
    dtype: The expected `dtype` of `grid` and `probs`.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
     quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
      representing the sample points and the corresponding (possibly
      normalized) weight.

  Raises:
    ValueError: if `quadrature_grid_and_probs is not None` and
      `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])`
  """
  with tf.name_scope(name or "process_quadrature_grid_and_probs"):
    if quadrature_grid_and_probs is None:
      grid, probs = np.polynomial.hermite.hermgauss(deg=8)
      grid = grid.astype(dtype_util.as_numpy_dtype(dtype))
      probs = probs.astype(dtype_util.as_numpy_dtype(dtype))
      probs /= np.linalg.norm(probs, ord=1, keepdims=True)
      grid = tf.convert_to_tensor(value=grid, name="grid", dtype=dtype)
      probs = tf.convert_to_tensor(value=probs, name="probs", dtype=dtype)
      return grid, probs

    grid, probs = tuple(quadrature_grid_and_probs)
    grid = tf.convert_to_tensor(value=grid, name="grid", dtype=dtype)
    probs = tf.convert_to_tensor(
        value=probs, name="unnormalized_probs", dtype=dtype)
    probs /= tf.norm(tensor=probs, ord=1, axis=-1, keepdims=True, name="probs")

    def _static_event_size(x):
      """Returns the static size of a specific dimension or `None`."""
      return tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(x.shape, 1)[-1])

    m, n = _static_event_size(probs), _static_event_size(grid)
    if m is not None and n is not None:
      if m != n:
        raise ValueError("`quadrature_grid_and_probs` must be a `tuple` of "
                         "same-length zero-th-dimension `Tensor`s "
                         "(saw lengths {}, {})".format(m, n))
    elif validate_args:
      assertions = [
          assert_util.assert_equal(
              dimension_size(probs, axis=-1),
              dimension_size(grid, axis=-1),
              message=("`quadrature_grid_and_probs` must be a `tuple` of "
                       "same-length zero-th-dimension `Tensor`s")),
      ]
      with tf.control_dependencies(assertions):
        grid = tf.identity(grid)
        probs = tf.identity(probs)
    return grid, probs

Validates quadrature grid probs or computes them as necessary .


def parent_frame_arguments():
  """Returns parent frame arguments.

  When called inside a function, returns a dictionary with the caller's function
  arguments. These are positional arguments and keyword arguments (**kwargs),
  while variable arguments (*varargs) are excluded.

  When called at global scope, this will return an empty dictionary, since there
  are no arguments.

  WARNING: If caller function argument names are overloaded before invoking
  this method, then values will reflect the overloaded value. For this reason,
  we recommend calling `parent_frame_arguments` at the beginning of the
  function.
  """
  # All arguments and the names used for *varargs, and **kwargs
  arg_names, variable_arg_name, keyword_arg_name, local_vars = (
      tf_inspect._inspect.getargvalues(  # pylint: disable=protected-access
          # Get the first frame of the caller of this method.
          tf_inspect._inspect.stack()[1][0]))  # pylint: disable=protected-access

  # Remove the *varargs, and flatten the **kwargs. Both are
  # nested lists.
  local_vars.pop(variable_arg_name, {})
  keyword_args = local_vars.pop(keyword_arg_name, {})

  final_args = {}
  # Copy over arguments and their values. In general, local_vars
  # may contain more than just the arguments, since this method
  # can be called anywhere in a function.
  for arg_name in arg_names:
    final_args[arg_name] = local_vars.pop(arg_name)
  final_args.update(keyword_args)

  return final_args


def expand_to_vector(x, tensor_name=None, op_name=None, validate_args=False):
  """Transform a 0-D or 1-D `Tensor` to be 1-D.

  For user convenience, many parts of the TensorFlow Probability API accept
  inputs of rank 0 or 1 -- i.e., allowing an `event_shape` of `[5]` to be passed
  to the API as either `5` or `[5]`.  This function can be used to transform
  such an argument to always be 1-D.

  NOTE: Python or NumPy values will be converted to `Tensor`s with standard type
  inference/conversion.  In particular, an empty list or tuple will become an
  empty `Tensor` with dtype `float32`.  Callers should convert values to
  `Tensor`s before calling this function if different behavior is desired
  (e.g. converting empty lists / other values to `Tensor`s with dtype `int32`).

  Args:
    x: A 0-D or 1-D `Tensor`.
    tensor_name: Python `str` name for `Tensor`s created by this function.
    op_name: Python `str` name for `Op`s created by this function.
    validate_args: Python `bool, default `False`.  When `True`, arguments may be
      checked for validity at execution time, possibly degrading runtime
      performance.  When `False`, invalid inputs may silently render incorrect
        outputs.
  Returns:
    vector: a 1-D `Tensor`.
  """
  with tf.name_scope(op_name or "expand_to_vector"):
    x = tf.convert_to_tensor(value=x, name="x")
    ndims = tensorshape_util.rank(x.shape)

    if ndims is None:
      # Maybe expand ndims from 0 to 1.
      if validate_args:
        x = with_dependencies([
            assert_util.assert_rank_at_most(
                x, 1, message="Input is neither scalar nor vector.")
        ], x)
      ndims = tf.rank(x)
      expanded_shape = pick_vector(
          tf.equal(ndims, 0), np.array([1], dtype=np.int32), tf.shape(input=x))
      return tf.reshape(x, expanded_shape)

    elif ndims == 0:
      # Definitely expand ndims from 0 to 1.
      x_const = tf.get_static_value(x)
      if x_const is not None:
        return tf.convert_to_tensor(
            value=dtype_util.as_numpy_dtype(x.dtype)([x_const]),
            name=tensor_name)

      else:
        return tf.reshape(x, [1])

    elif ndims != 1:
      raise ValueError("Input is neither scalar nor vector.")

    # ndims == 1
    return x

def
expand_to_vector
(
x
,
tensor_name
=
,
op_name
=
,
validate_args
=
False
)
:
with
tf
.
name_scope
(
op_name
or
"expand_to_vector"
)
:
x
=
tf
.
convert_to_tensor
(
value
=
x
,
name
=
"x"
)
ndims
=
tensorshape_util
.
rank
(
x
.
shape
)
if
ndims
is
:
# Maybe expand ndims from 0 to 1.
if
validate_args
:
x
=
with_dependencies
(
[
assert_util
.
assert_rank_at_most
(
x
,
1
,
message
=
"Input is neither scalar nor vector."
)
]
,
x
)
ndims
=
tf
.
rank
(
x
)
expanded_shape
=
pick_vector
(
tf
.
equal
(
ndims
,
0
)
,
np
.
array
(
[
1
]
,
dtype
=
np
.
int32
)
,
tf
.
shape
(
input
=
x
)
)
return
tf
.
reshape
(
x
,
expanded_shape
)
elif
ndims
==
0
:
# Definitely expand ndims from 0 to 1.
x_const
=
tf
.
get_static_value
(
x
)
if
x_const
is
not
:
return
tf
.
convert_to_tensor
(
value
=
dtype_util
.
as_numpy_dtype
(
x
.
dtype
)
(
[
x_const
]
)
,
name
=
tensor_name
)
else
:
return
tf
.
reshape
(
x
,
[
1
]
)
elif
ndims
!=
1
:
raise
ValueError
(
"Input is neither scalar nor vector."
)
# ndims == 1
return
x

Transform a 0 - D or 1 - D Tensor to be 1 - D .


def with_dependencies(dependencies, output_tensor, name=None):
  """Produces the content of `output_tensor` only after `dependencies`.

  In some cases, a user may want the output of an operation to be consumed
  externally only after some other dependencies have run first. This function
  returns `output_tensor`, but only after all operations in `dependencies` have
  run. Note that this means that there is no guarantee that `output_tensor` will
  be evaluated after any `dependencies` have run.

  See also `tf.tuple` and `tf.group`.

  Args:
    dependencies: Iterable of operations to run before this op finishes.
    output_tensor: A `Tensor` or `IndexedSlices` that will be returned.
    name: (Optional) A name for this operation.

  Returns:
    output_with_deps: Same as `output_tensor` but with embedded dependencies.

  Raises:
    TypeError: if `output_tensor` is not a `Tensor` or `IndexedSlices`.
  """
  if tf.executing_eagerly():
    return output_tensor
  with tf.name_scope(name or "control_dependency") as name:
    with tf.control_dependencies(d for d in dependencies if d is not None):
      output_tensor = tf.convert_to_tensor(value=output_tensor)
      if isinstance(output_tensor, tf.Tensor):
        return tf.identity(output_tensor, name=name)
      else:
        return tf.IndexedSlices(
            tf.identity(output_tensor.values, name=name),
            output_tensor.indices,
            output_tensor.dense_shape)

def with_dependencies ( dependencies , output_tensor , name = None ) : if tf . executing_eagerly ( ) : return output_tensor with tf . name_scope ( name or "control_dependency" ) as name : with tf . control_dependencies ( d for d in dependencies if d is not None ) : output_tensor = tf . convert_to_tensor ( value = output_tensor ) if isinstance ( output_tensor , tf . Tensor ) : return tf . identity ( output_tensor , name = name ) else : return tf . IndexedSlices ( tf . identity ( output_tensor . values , name = name ) , output_tensor . indices , output_tensor . dense_shape )

Produces the content of output_tensor only after dependencies .


def _maybe_validate_rightmost_transposed_ndims(
    rightmost_transposed_ndims, validate_args, name=None):
  """Checks that `rightmost_transposed_ndims` is valid."""
  with tf.name_scope(name or 'maybe_validate_rightmost_transposed_ndims'):
    assertions = []
    if not dtype_util.is_integer(rightmost_transposed_ndims.dtype):
      raise TypeError('`rightmost_transposed_ndims` must be integer type.')

    if tensorshape_util.rank(rightmost_transposed_ndims.shape) is not None:
      if tensorshape_util.rank(rightmost_transposed_ndims.shape) != 0:
        raise ValueError('`rightmost_transposed_ndims` must be a scalar, '
                         'saw rank: {}.'.format(
                             tensorshape_util.rank(
                                 rightmost_transposed_ndims.shape)))
    elif validate_args:
      assertions += [assert_util.assert_rank(rightmost_transposed_ndims, 0)]

    rightmost_transposed_ndims_ = tf.get_static_value(
        rightmost_transposed_ndims)
    msg = '`rightmost_transposed_ndims` must be non-negative.'
    if rightmost_transposed_ndims_ is not None:
      if rightmost_transposed_ndims_ < 0:
        raise ValueError(msg[:-1] + ', saw: {}.'.format(
            rightmost_transposed_ndims_))
    elif validate_args:
      assertions += [
          assert_util.assert_non_negative(
              rightmost_transposed_ndims, message=msg)
      ]

    return assertions

Checks that rightmost_transposed_ndims is valid .


def _maybe_validate_perm(perm, validate_args, name=None):
  """Checks that `perm` is valid."""
  with tf.name_scope(name or 'maybe_validate_perm'):
    assertions = []
    if not dtype_util.is_integer(perm.dtype):
      raise TypeError('`perm` must be integer type')

    msg = '`perm` must be a vector.'
    if tensorshape_util.rank(perm.shape) is not None:
      if tensorshape_util.rank(perm.shape) != 1:
        raise ValueError(
            msg[:-1] +
            ', saw rank: {}.'.format(tensorshape_util.rank(perm.shape)))
    elif validate_args:
      assertions += [assert_util.assert_rank(perm, 1, message=msg)]

    perm_ = tf.get_static_value(perm)
    msg = '`perm` must be a valid permutation vector.'
    if perm_ is not None:
      if not np.all(np.arange(np.size(perm_)) == np.sort(perm_)):
        raise ValueError(msg[:-1] + ', saw: {}.'.format(perm_))
    elif validate_args:
      assertions += [
          assert_util.assert_equal(
              tf.sort(perm), tf.range(tf.size(input=perm)), message=msg)
      ]

    return assertions


def _event_shape(self, shape, static_perm_to_shape):
    """Helper for _forward and _inverse_event_shape."""
    rightmost_ = tf.get_static_value(self.rightmost_transposed_ndims)
    if tensorshape_util.rank(shape) is None or rightmost_ is None:
      return tf.TensorShape(None)
    if tensorshape_util.rank(shape) < rightmost_:
      raise ValueError('Invalid shape: min event ndims={} but got {}'.format(
          rightmost_, shape))
    perm_ = tf.get_static_value(self.perm, partial=True)
    if perm_ is None:
      return shape[:tensorshape_util.rank(shape) - rightmost_].concatenate(
          [None] * int(rightmost_))
    # We can use elimination to reidentify a single None dimension.
    if sum(p is None for p in perm_) == 1:
      present = np.argsort([-1 if p is None else p for p in perm_])
      for i, p in enumerate(present[1:]):  # The -1 sorts to position 0.
        if i != p:
          perm_ = [i if p is None else p for p in perm_]
          break
    return shape[:tensorshape_util.rank(shape) - rightmost_].concatenate(
        static_perm_to_shape(shape[tensorshape_util.rank(shape) - rightmost_:],
                             perm_))

Helper for _forward and _inverse_event_shape .


def concatenate(x, other):
  """Returns the concatenation of the dimension in `x` and `other`.

  *Note:* If either `x` or `other` is completely unknown, concatenation will
  discard information about the other shape. In future, we might support
  concatenation that preserves this information for use with slicing.

  For more details, see `help(tf.TensorShape.concatenate)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    other: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    new_shape: an object like `x` whose elements are the concatenation of the
      dimensions in `x` and `other`.
  """
  return type(x)(tf.TensorShape(x).concatenate(other))

def concatenate ( x , other ) : return type ( x ) ( tf . TensorShape ( x ) . concatenate ( other ) )

Returns the concatenation of the dimension in x and other .


def dims(x):
  """Returns a list of dimension sizes, or `None` if `rank` is unknown.

  For more details, see `help(tf.TensorShape.dims)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    shape_as_list: list of sizes or `None` values representing each
      dimensions size if known. A size is `tf.Dimension` if input is a
      `tf.TensorShape` and an `int` otherwise.
  """
  if isinstance(x, tf.TensorShape):
    return x.dims
  r = tf.TensorShape(x).dims
  return None if r is None else list(map(tf.compat.dimension_value, r))

def dims ( x ) : if isinstance ( x , tf . TensorShape ) : return x . dims r = tf . TensorShape ( x ) . dims return None if r is None else list ( map ( tf . compat . dimension_value , r ) )

Returns a list of dimension sizes or None if rank is unknown .


def merge_with(x, other):
  """Returns a shape combining the information in `x` and `other`.

  The dimensions in `x` and `other` are merged elementwise, according to the
  rules defined for `tf.Dimension.merge_with()`.

  For more details, see `help(tf.TensorShape.merge_with)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    other: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    merged_shape: shape having `type(x)` containing the combined information of
      `x` and `other`.

  Raises:
    ValueError: If `x` and `other` are not compatible.
  """
  return type(x)(tf.TensorShape(x).merge_with(other))

def merge_with ( x , other ) : return type ( x ) ( tf . TensorShape ( x ) . merge_with ( other ) )

Returns a shape combining the information in x and other .


def with_rank_at_least(x, rank):  # pylint: disable=redefined-outer-name
  """Returns a shape based on `x` with at least the given `rank`.

  For more details, see `help(tf.TensorShape.with_rank_at_least)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    rank: An `int` representing the minimum rank of `x` or else an assertion is
      raised.

  Returns:
    shape: a shape having `type(x)` but guaranteed to have at least the given
      rank (or else an assertion was raised).

  Raises:
    ValueError: If `x` does not represent a shape with at least the given
      `rank`.
  """
  return type(x)(tf.TensorShape(x).with_rank_at_least(rank))

def
with_rank_at_least
(
x
,
rank
)
:
# pylint: disable=redefined-outer-name
return
type
(
x
)
(
tf
.
TensorShape
(
x
)
.
with_rank_at_least
(
rank
)
)

Returns a shape based on x with at least the given rank .


def _check_equal_shape(name,
                       static_shape,
                       dynamic_shape,
                       static_target_shape,
                       dynamic_target_shape=None):
  """Check that source and target shape match, statically if possible."""

  static_target_shape = tf.TensorShape(static_target_shape)
  if tensorshape_util.is_fully_defined(
      static_shape) and tensorshape_util.is_fully_defined(static_target_shape):
    if static_shape != static_target_shape:
      raise ValueError("{}: required shape {} but found {}".
                       format(name, static_target_shape, static_shape))
    return None
  else:
    if dynamic_target_shape is None:
      if tensorshape_util.is_fully_defined(static_target_shape):
        dynamic_target_shape = tensorshape_util.as_list(static_target_shape)
      else:
        raise ValueError("{}: cannot infer target shape: no dynamic shape "
                         "specified and static shape {} is not fully defined".
                         format(name, static_target_shape))
    return assert_util.assert_equal(
        dynamic_shape,
        dynamic_target_shape,
        message=("{}: required shape {}".format(name, static_target_shape)))

Check that source and target shape match statically if possible .


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 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

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

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

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()

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 .