<EOL>batch_shape = tf.TensorShape(batch_shape).as_list()<EOL>slices = []<EOL>batch_rank = len(batch_shape)<EOL>arbitrary_slices = hps.tuples(<EOL>hps.one_of(hps.just(None), hps.integers(min_value=-<NUM_LIT:100>, max_value=<NUM_LIT:100>)),<EOL>hps.one_of(hps.just(None), hps.integers(min_value=-<NUM_LIT:100>, max_value=<NUM_LIT:100>)),<EOL>hps.one_of(<EOL>hps.just(None),<EOL>hps.integers(min_value=-<NUM_LIT:100>, max_value=<NUM_LIT:100>).filter(lambda x: x != <NUM_LIT:0>))<EOL>).map(lambda tup: slice(*tup))<EOL>nslc_before_ellipsis = draw(hps.integers(min_value=<NUM_LIT:0>, max_value=batch_rank))<EOL>for i in range(nslc_before_ellipsis):<EOL><INDENT>slc = draw(<EOL>hps.one_of(<EOL>hps.integers(min_value=<NUM_LIT:0>, max_value=batch_shape[i] - <NUM_LIT:1>),<EOL>arbitrary_slices))<EOL>slices.append(slc)<EOL><DEDENT>has_ellipsis = draw(hps.booleans().map(lambda x: (Ellipsis, x)))[<NUM_LIT:1>]<EOL>nslc_after_ellipsis = draw(<EOL>hps.integers(min_value=<NUM_LIT:0>, max_value=batch_rank - nslc_before_ellipsis))<EOL>if has_ellipsis:<EOL><INDENT>slices.append(Ellipsis)<EOL>remain_start, remain_end = (batch_rank - nslc_after_ellipsis, batch_rank)<EOL><DEDENT>else:<EOL><INDENT>remain_start = nslc_before_ellipsis<EOL>remain_end = nslc_before_ellipsis + nslc_after_ellipsis<EOL><DEDENT>for i in range(remain_start, remain_end):<EOL><INDENT>slc = draw(<EOL>hps.one_of(<EOL>hps.integers(min_value=<NUM_LIT:0>, max_value=batch_shape[i] - <NUM_LIT:1>),<EOL>arbitrary_slices))<EOL>slices.append(slc)<EOL><DEDENT>newaxis_positions = draw(<EOL>hps.lists(hps.integers(min_value=<NUM_LIT:0>, max_value=len(slices)), max_size=<NUM_LIT:4>))<EOL>for i in sorted(newaxis_positions, reverse=True):<EOL><INDENT>slices.insert(i, tf.newaxis)<EOL><DEDENT>slices = tuple(slices)<EOL>if len(slices) == <NUM_LIT:1> and draw(hps.booleans()):<EOL><INDENT>return slices[<NUM_LIT:0>]<EOL><DEDENT>return slices<EOL>

Samples a legal (possibly empty) slice for shape batch_shape.


pretty_slices = []<EOL>slices = slices if isinstance(slices, tuple) else (slices,)<EOL>for slc in slices:<EOL><INDENT>if slc == Ellipsis:<EOL><INDENT>pretty_slices.append('<STR_LIT>')<EOL><DEDENT>elif isinstance(slc, slice):<EOL><INDENT>pretty_slices.append('<STR_LIT>'.format(<EOL>*['<STR_LIT>' if s is None else s for s in (slc.start, slc.stop, slc.step)]))<EOL><DEDENT>elif isinstance(slc, int) or tf.is_tensor(slc):<EOL><INDENT>pretty_slices.append(str(slc))<EOL><DEDENT>elif slc is tf.newaxis:<EOL><INDENT>pretty_slices.append('<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(type(slc)))<EOL><DEDENT><DEDENT>return pretty_slices<EOL>

Returns a list of strings describing the items in `slices`.


<EOL>float32s = hps.floats(-<NUM_LIT:200>, <NUM_LIT:200>, allow_nan=False, allow_infinity=False)<EOL>def mapper(x):<EOL><INDENT>result = assert_util.assert_finite(<EOL>constraint_fn(tf.convert_to_tensor(value=x)),<EOL>message='<STR_LIT>')<EOL>if tf.executing_eagerly():<EOL><INDENT>return result.numpy()<EOL><DEDENT>return result<EOL><DEDENT>return hpnp.arrays(<EOL>dtype=np.float32, shape=param_shape, elements=float32s).map(mapper)<EOL>

Draws the value of a single distribution parameter.


def broadcasting_params(draw, dist_name, batch_shape, event_dim=None):

_, params_event_ndims = INSTANTIABLE_DISTS[dist_name]<EOL>if event_dim is None:<EOL><INDENT>event_dim = draw(hps.integers(min_value=<NUM_LIT:2>, max_value=<NUM_LIT:6>))<EOL><DEDENT>remaining_params = set(params_event_ndims.keys())<EOL>params_to_use = []<EOL>while remaining_params:<EOL><INDENT>param = draw(hps.one_of(map(hps.just, remaining_params)))<EOL>params_to_use.append(param)<EOL>remaining_params.remove(param)<EOL>for mutex_set in MUTEX_PARAMS:<EOL><INDENT>if param in mutex_set:<EOL><INDENT>remaining_params -= mutex_set<EOL><DEDENT><DEDENT><DEDENT>param_batch_shapes = draw(broadcasting_shapes(batch_shape, params_to_use))<EOL>params_kwargs = dict()<EOL>for param in params_to_use:<EOL><INDENT>param_batch_shape = param_batch_shapes[param]<EOL>param_event_rank = params_event_ndims[param]<EOL>params_kwargs[param] = tf.convert_to_tensor(<EOL>value=draw(<EOL>single_param(<EOL>constraint_for(dist_name, param),<EOL>(tensorshape_util.as_list(param_batch_shape) +<EOL>[event_dim] * param_event_rank))),<EOL>dtype=tf.float32)<EOL><DEDENT>return params_kwargs<EOL>

Draws a dict of parameters which should yield the given batch shape.


if dist_name is None:<EOL><INDENT>def dist_filter(dist_name):<EOL><INDENT>if not eligibility_filter(dist_name):<EOL><INDENT>return False<EOL><DEDENT>dist_cls, _ = INSTANTIABLE_DISTS[dist_name]<EOL>if (tf.executing_eagerly() and<EOL>issubclass(dist_cls, tfd.TransformedDistribution)):<EOL><INDENT>return False<EOL><DEDENT>return True<EOL><DEDENT>dist_name = draw(<EOL>hps.one_of(<EOL>map(hps.just,<EOL>[k for k in INSTANTIABLE_DISTS.keys() if dist_filter(k)])))<EOL><DEDENT>dist_cls, _ = INSTANTIABLE_DISTS[dist_name]<EOL>if dist_name == '<STR_LIT>':<EOL><INDENT>return independents(draw, batch_shape, event_dim)<EOL><DEDENT>if dist_name == '<STR_LIT>':<EOL><INDENT>return mixtures_same_family(draw, batch_shape, event_dim)<EOL><DEDENT>if dist_name == '<STR_LIT>':<EOL><INDENT>return transformed_distributions(draw, batch_shape, event_dim)<EOL><DEDENT>if batch_shape is None:<EOL><INDENT>batch_shape = draw(batch_shapes())<EOL><DEDENT>params_kwargs = broadcasting_params(<EOL>draw, dist_name, batch_shape, event_dim=event_dim)<EOL>params_constrained = constraint_for(dist_name)(params_kwargs)<EOL>assert_shapes_unchanged(params_kwargs, params_constrained)<EOL>params_constrained['<STR_LIT>'] = True<EOL>return dist_cls(**params_constrained), batch_shape<EOL>

Samples one a set of supported distributions.


new_high = tf.maximum(low + tf.abs(low) * <NUM_LIT> + <NUM_LIT>, high)<EOL>reduce_dims = []<EOL>if (tensorshape_util.rank(new_high.shape) ><EOL>tensorshape_util.rank(high.shape)):<EOL><INDENT>reduced_leading_axes = tf.range(<EOL>tensorshape_util.rank(new_high.shape) -<EOL>tensorshape_util.rank(high.shape))<EOL>new_high = tf.math.reduce_max(<EOL>input_tensor=new_high, axis=reduced_leading_axes)<EOL><DEDENT>reduce_dims = [<EOL>d for d in range(tensorshape_util.rank(high.shape))<EOL>if high.shape[d] < new_high.shape[d]<EOL>]<EOL>if reduce_dims:<EOL><INDENT>new_high = tf.math.reduce_max(<EOL>input_tensor=new_high, axis=reduce_dims, keepdims=True)<EOL><DEDENT>return new_high<EOL>

Returns a value with shape matching `high` and gt broadcastable `low`.


@kullback_leibler.RegisterKL(Laplace, Laplace)<EOL>def _kl_laplace_laplace(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>distance = tf.abs(a.loc - b.loc)<EOL>ratio = a.scale / b.scale<EOL>return (-tf.math.log(ratio) - <NUM_LIT:1> + distance / b.scale +<EOL>ratio * tf.exp(-distance / a.scale))<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b Laplace.

    Args:
      a: instance of a Laplace distribution object.
      b: instance of a Laplace distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_laplace_laplace".

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], tf.float32)<EOL>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(scale)] if validate_args else []):<EOL><INDENT>self._loc = tf.identity(loc)<EOL>self._scale = tf.identity(scale)<EOL>dtype_util.assert_same_float_dtype([self._loc, self._scale])<EOL><DEDENT>super(Laplace, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._loc, self._scale],<EOL>name=name)<EOL><DEDENT>

Construct Laplace distribution with parameters `loc` and `scale`.

        The parameters `loc` and `scale` must be shaped in a way that supports
        broadcasting (e.g., `loc / scale` is a valid operation).

        Args:
          loc: Floating point tensor which characterizes the location (center)
            of the distribution.
          scale: Positive floating point tensor which characterizes the spread of
            the distribution.
          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.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          TypeError: if `loc` and `scale` are of different dtype.


return self._loc<EOL>

Distribution parameter for the location.


return self._scale<EOL>


def _maybe_validate_distributions(distributions, dtype_override, validate_args):

assertions = []<EOL>if not _is_iterable(distributions) or not distributions:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if dtype_override is None:<EOL><INDENT>dts = [<EOL>dtype_util.base_dtype(d.dtype)<EOL>for d in distributions<EOL>if d.dtype is not None<EOL>]<EOL>if dts[<NUM_LIT:1>:] != dts[:-<NUM_LIT:1>]:<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>set(dtype_util.name(dt) for dt in dts)))<EOL><DEDENT><DEDENT>for d in distributions:<EOL><INDENT>if tensorshape_util.rank(d.event_shape) is not None:<EOL><INDENT>if tensorshape_util.rank(d.event_shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>tensorshape_util.rank(d.event_shape)))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_equal(<EOL><NUM_LIT:1>, tf.size(input=d.event_shape_tensor()),<EOL>message='<STR_LIT>'))<EOL><DEDENT><DEDENT>batch_shapes = [d.batch_shape for d in distributions]<EOL>if all(tensorshape_util.is_fully_defined(b) for b in batch_shapes):<EOL><INDENT>if batch_shapes[<NUM_LIT:1>:] != batch_shapes[:-<NUM_LIT:1>]:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(batch_shapes))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>batch_shapes = [<EOL>tensorshape_util.as_list(d.batch_shape)  <EOL>if tensorshape_util.is_fully_defined(d.batch_shape) else<EOL>d.batch_shape_tensor() for d in distributions<EOL>]<EOL>assertions.extend(<EOL>assert_util.assert_equal(  <EOL>b1, b2,<EOL>message='<STR_LIT>')<EOL>for b1, b2 in zip(batch_shapes[<NUM_LIT:1>:], batch_shapes[:-<NUM_LIT:1>]))<EOL><DEDENT>return assertions<EOL>

Checks that `distributions` satisfies all assumptions.


@kullback_leibler.RegisterKL(Blockwise, Blockwise)<EOL>def _kl_blockwise_blockwise(b0, b1, name=None):

if len(b0.distributions) != len(b1.distributions):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>b0_event_sizes = [_event_size(d) for d in b0.distributions]<EOL>b1_event_sizes = [_event_size(d) for d in b1.distributions]<EOL>assertions = []<EOL>message = ('<STR_LIT>'<EOL>'<STR_LIT>')<EOL>if (all(isinstance(event_size, int) for event_size in b0_event_sizes) and<EOL>all(isinstance(event_size, int) for event_size in b1_event_sizes)):<EOL><INDENT>if b0_event_sizes != b1_event_sizes:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>if b0.validate_args or b1.validate_args:<EOL><INDENT>assertions.extend(<EOL>assert_util.assert_equal(  <EOL>e1, e2, message=message)<EOL>for e1, e2 in zip(b0_event_sizes, b1_event_sizes))<EOL><DEDENT><DEDENT>with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>return sum([<EOL>kullback_leibler.kl_divergence(d1, d2) for d1, d2 in zip(<EOL>b0.distributions, b1.distributions)])<EOL><DEDENT><DEDENT>

Calculate the batched KL divergence KL(b0 || b1) with b0 and b1 Blockwise distributions.

    Args:
      b0: instance of a Blockwise distribution object.
      b1: instance of a Blockwise distribution object.
      name: (optional) Name to use for created operations. Default is
        "kl_blockwise_blockwise".

    Returns:
      kl_blockwise_blockwise: `Tensor`. The batchwise KL(b0 || b1).


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._assertions = _maybe_validate_distributions(<EOL>distributions, dtype_override, validate_args)<EOL>if dtype_override is not None:<EOL><INDENT>dtype = dtype_override<EOL><DEDENT>else:<EOL><INDENT>dtype = set(<EOL>dtype_util.base_dtype(d.dtype)<EOL>for d in distributions<EOL>if d.dtype is not None)<EOL>if len(dtype) == <NUM_LIT:0>:  <EOL><INDENT>dtype = tf.float32<EOL><DEDENT>elif len(dtype) == <NUM_LIT:1>:<EOL><INDENT>dtype = dtype.pop()<EOL><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>reparameterization_type = set(d.reparameterization_type<EOL>for d in distributions)<EOL>reparameterization_type = (reparameterization_type.pop()<EOL>if len(reparameterization_type) == <NUM_LIT:1><EOL>else reparameterization.NOT_REPARAMETERIZED)<EOL>self._distributions = distributions<EOL>super(Blockwise, self).__init__(<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization_type,<EOL>parameters=parameters,<EOL>graph_parents=_flatten(d._graph_parents for d in distributions),  <EOL>name=name)<EOL><DEDENT>

Construct the `Blockwise` distribution.

        Args:
          distributions: Python `list` of `tfp.distributions.Distribution`
            instances. All distribution instances must have the same `batch_shape`
            and all must have `event_ndims==1`, i.e., be vector-variate
            distributions.
          dtype_override: samples of `distributions` will be cast to this `dtype`.
            If unspecified, all `distributions` must have the same `dtype`.
            Default value: `None` (i.e., do not cast).
          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.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


static_target_shape = tf.TensorShape(static_target_shape)<EOL>if tensorshape_util.is_fully_defined(<EOL>static_shape) and tensorshape_util.is_fully_defined(static_target_shape):<EOL><INDENT>if static_shape != static_target_shape:<EOL><INDENT>raise ValueError("<STR_LIT>".<EOL>format(name, static_target_shape, static_shape))<EOL><DEDENT>return None<EOL><DEDENT>else:<EOL><INDENT>if dynamic_target_shape is None:<EOL><INDENT>if tensorshape_util.is_fully_defined(static_target_shape):<EOL><INDENT>dynamic_target_shape = tensorshape_util.as_list(static_target_shape)<EOL><DEDENT>else:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".<EOL>format(name, static_target_shape))<EOL><DEDENT><DEDENT>return assert_util.assert_equal(<EOL>dynamic_shape,<EOL>dynamic_target_shape,<EOL>message=("<STR_LIT>".format(name, static_target_shape)))<EOL><DEDENT>

Check that source and target shape match, statically if possible.


full_ndims = distribution_util.prefer_static_shape(<EOL>full_sample_and_batch_shape)[<NUM_LIT:0>]<EOL>partial_batch_ndims = (<EOL>tensorshape_util.rank(partial_batch_dist.batch_shape)  <EOL>if tensorshape_util.rank(partial_batch_dist.batch_shape) is not None<EOL>else distribution_util.prefer_static_shape(<EOL>partial_batch_dist.batch_shape_tensor())[<NUM_LIT:0>])<EOL>num_broadcast_dims = full_ndims - partial_batch_ndims<EOL>expected_partial_batch_shape = (<EOL>full_sample_and_batch_shape[num_broadcast_dims:])<EOL>expected_partial_batch_shape_static = tf.get_static_value(<EOL>full_sample_and_batch_shape[num_broadcast_dims:])<EOL>num_broadcast_dims_static = tf.get_static_value(num_broadcast_dims)<EOL>if num_broadcast_dims_static is not None:<EOL><INDENT>if num_broadcast_dims_static < <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>.format(partial_batch_dist))<EOL><DEDENT><DEDENT>if (expected_partial_batch_shape_static is not None and<EOL>tensorshape_util.is_fully_defined(partial_batch_dist.batch_shape)):<EOL><INDENT>if (partial_batch_dist.batch_shape and<EOL>any(expected_partial_batch_shape_static != tensorshape_util.as_list(<EOL>partial_batch_dist.batch_shape))):<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>expected_partial_batch_shape_static,<EOL>partial_batch_dist.batch_shape<EOL>))<EOL><DEDENT><DEDENT>runtime_assertions = []<EOL>if validate_args:<EOL><INDENT>runtime_assertions.append(<EOL>assert_util.assert_greater_equal(<EOL>tf.convert_to_tensor(value=num_broadcast_dims, dtype=tf.int32),<EOL>tf.zeros((), dtype=tf.int32),<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>partial_batch_dist))))<EOL>runtime_assertions.append(<EOL>assert_util.assert_equal(<EOL>expected_partial_batch_shape,<EOL>partial_batch_dist.batch_shape_tensor(),<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>"),<EOL>name="<STR_LIT>"))<EOL><DEDENT>with tf.control_dependencies(runtime_assertions):<EOL><INDENT>return full_sample_and_batch_shape[:num_broadcast_dims]<EOL><DEDENT>

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


def backward_pass_step(state,<EOL>filtered_parameters):<EOL><INDENT>"""<STR_LIT>"""<EOL>(filtered_mean, filtered_cov,<EOL>predicted_mean, predicted_cov) = filtered_parameters<EOL>transition_matrix = get_transition_matrix_for_timestep(state.timestep)<EOL>next_posterior_mean = state.backward_mean<EOL>next_posterior_cov = state.backward_cov<EOL>posterior_mean, posterior_cov = backward_smoothing_update(<EOL>filtered_mean,<EOL>filtered_cov,<EOL>predicted_mean,<EOL>predicted_cov,<EOL>next_posterior_mean,<EOL>next_posterior_cov,<EOL>transition_matrix)<EOL>return BackwardPassState(backward_mean=posterior_mean,<EOL>backward_cov=posterior_cov,<EOL>timestep=state.timestep-<NUM_LIT:1>)<EOL><DEDENT>return backward_pass_step<EOL>

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


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

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


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

Build a callable that performs one step of Kalman filtering.

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

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


def linear_gaussian_update(<EOL>prior_mean, prior_cov, observation_matrix, observation_noise, x_observed):

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

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.


def kalman_transition(filtered_mean, filtered_cov,<EOL>transition_matrix, transition_noise):

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

Propagate a filtered distribution through a transition model.


def mean_step(previous_means, t):<EOL><INDENT>"""<STR_LIT>"""<EOL>previous_latent_mean, _ = previous_means<EOL>latent_mean = _propagate_mean(previous_latent_mean,<EOL>get_transition_matrix_for_timestep(t - <NUM_LIT:1>),<EOL>get_transition_noise_for_timestep(t - <NUM_LIT:1>))<EOL>observation_mean = _propagate_mean(latent_mean,<EOL>get_observation_matrix_for_timestep(t),<EOL>get_observation_noise_for_timestep(t))<EOL>return (latent_mean, observation_mean)<EOL><DEDENT>return mean_step<EOL>

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 cov_step(previous_covs, t):<EOL><INDENT>"""<STR_LIT>"""<EOL>previous_latent_cov, _ = previous_covs<EOL>latent_cov = _propagate_cov(<EOL>previous_latent_cov,<EOL>get_transition_matrix_for_timestep(t - <NUM_LIT:1>),<EOL>get_transition_noise_for_timestep(t - <NUM_LIT:1>))<EOL>observation_cov = _propagate_cov(<EOL>latent_cov,<EOL>get_observation_matrix_for_timestep(t),<EOL>get_observation_noise_for_timestep(t))<EOL>return (latent_cov, observation_cov)<EOL><DEDENT>return cov_step<EOL>

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 sample_step(sampled_prev, t):<EOL><INDENT>"""<STR_LIT>"""<EOL>latent_prev, _ = sampled_prev<EOL>transition_matrix = get_transition_matrix_for_timestep(t - <NUM_LIT:1>)<EOL>transition_noise = get_transition_noise_for_timestep(t - <NUM_LIT:1>)<EOL>latent_pred = transition_matrix.matmul(latent_prev)<EOL>latent_sampled = latent_pred + transition_noise.sample(<EOL>sample_shape=_augment_sample_shape(<EOL>transition_noise,<EOL>full_sample_and_batch_shape,<EOL>validate_args),<EOL>seed=stream())[..., tf.newaxis]<EOL>observation_matrix = get_observation_matrix_for_timestep(t)<EOL>observation_noise = get_observation_noise_for_timestep(t)<EOL>observation_pred = observation_matrix.matmul(latent_sampled)<EOL>observation_sampled = observation_pred + observation_noise.sample(<EOL>sample_shape=_augment_sample_shape(<EOL>observation_noise,<EOL>full_sample_and_batch_shape,<EOL>validate_args),<EOL>seed=stream())[..., tf.newaxis]<EOL>return (latent_sampled, observation_sampled)<EOL><DEDENT>return sample_step<EOL>

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 pushforward_latents_step(_, latent_t_mean_cov):<EOL><INDENT>"""<STR_LIT>"""<EOL>t, latent_mean, latent_cov = latent_t_mean_cov<EOL>observation_matrix = get_observation_matrix_for_timestep(t)<EOL>observation_noise = get_observation_noise_for_timestep(t)<EOL>observation_mean = _propagate_mean(latent_mean,<EOL>observation_matrix,<EOL>observation_noise)<EOL>observation_cov = _propagate_cov(latent_cov,<EOL>observation_matrix,<EOL>observation_noise)<EOL>return (observation_mean, observation_cov)<EOL><DEDENT>return pushforward_latents_step<EOL>

Build a callable to push latent means/covs to observed means/covs.

    Args:
      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:
      pushforward_latents_step: a callable that computes the observation mean and
      covariance at time `t`, given latent mean and covariance at time `t`.


return linop.matmul(mean) + dist.mean()[..., tf.newaxis]<EOL>

Propagate a mean through linear Gaussian transformation.


<EOL>return linop.matmul(linop.matmul(cov), adjoint_arg=True) + dist.covariance()<EOL>

Propagate covariance through linear Gaussian transformation.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self.num_timesteps = tf.convert_to_tensor(<EOL>value=num_timesteps, name="<STR_LIT>")<EOL>self.initial_state_prior = initial_state_prior<EOL>self.initial_step = initial_step<EOL>self.final_step = self.initial_step + self.num_timesteps<EOL>dtype = initial_state_prior.dtype<EOL>def _maybe_make_linop(x, is_square=None, name=None):<EOL><INDENT>"""<STR_LIT>"""<EOL>if not isinstance(x, tfl.LinearOperator):<EOL><INDENT>x = tfl.LinearOperatorFullMatrix(<EOL>tf.convert_to_tensor(value=x, dtype=dtype),<EOL>is_square=is_square,<EOL>name=name)<EOL><DEDENT>return x<EOL><DEDENT>def _maybe_make_callable_from_linop(f, name, make_square_linop=None):<EOL><INDENT>"""<STR_LIT>"""<EOL>if not callable(f):<EOL><INDENT>linop = _maybe_make_linop(f, is_square=make_square_linop, name=name)<EOL>f = lambda t: linop<EOL><DEDENT>return f<EOL><DEDENT>self.get_transition_matrix_for_timestep = (<EOL>_maybe_make_callable_from_linop(<EOL>transition_matrix,<EOL>name="<STR_LIT>",<EOL>make_square_linop=True))<EOL>self.get_observation_matrix_for_timestep = (<EOL>_maybe_make_callable_from_linop(<EOL>observation_matrix, name="<STR_LIT>"))<EOL>def _maybe_make_callable(f):<EOL><INDENT>if not callable(f):<EOL><INDENT>return lambda t: f<EOL><DEDENT>return f<EOL><DEDENT>self.get_transition_noise_for_timestep = _maybe_make_callable(<EOL>transition_noise)<EOL>self.get_observation_noise_for_timestep = _maybe_make_callable(<EOL>observation_noise)<EOL>self.latent_size = distribution_util.prefer_static_value(<EOL>initial_state_prior.event_shape_tensor())[-<NUM_LIT:1>]<EOL>self.observation_size = distribution_util.prefer_static_value(<EOL>self.get_observation_matrix_for_timestep(<EOL>self.initial_step).shape_tensor())[-<NUM_LIT:2>]<EOL>self.runtime_assertions = []<EOL>if validate_args:<EOL><INDENT>transition_matrix = (<EOL>self.get_transition_matrix_for_timestep(self.initial_step))<EOL>transition_noise = (<EOL>self.get_transition_noise_for_timestep(self.initial_step))<EOL>observation_matrix = (<EOL>self.get_observation_matrix_for_timestep(self.initial_step))<EOL>observation_noise = (<EOL>self.get_observation_noise_for_timestep(self.initial_step))<EOL>dtype_util.assert_same_float_dtype([<EOL>initial_state_prior, transition_matrix, transition_noise,<EOL>observation_matrix, observation_noise<EOL>])<EOL>latent_size_ = tf.get_static_value(self.latent_size)<EOL>observation_size_ = tf.get_static_value(<EOL>self.observation_size)<EOL>runtime_assertions = [<EOL>_check_equal_shape(<EOL>name="<STR_LIT>",<EOL>static_shape=transition_matrix.shape[-<NUM_LIT:2>:],<EOL>dynamic_shape=transition_matrix.shape_tensor()[-<NUM_LIT:2>:],<EOL>static_target_shape=[latent_size_, latent_size_],<EOL>dynamic_target_shape=[self.latent_size, self.latent_size]),<EOL>_check_equal_shape(<EOL>name="<STR_LIT>",<EOL>static_shape=observation_matrix.shape[-<NUM_LIT:2>:],<EOL>dynamic_shape=observation_matrix.shape_tensor()[-<NUM_LIT:2>:],<EOL>static_target_shape=[observation_size_, latent_size_],<EOL>dynamic_target_shape=[self.observation_size, self.latent_size]),<EOL>_check_equal_shape(<EOL>name="<STR_LIT>",<EOL>static_shape=initial_state_prior.event_shape,<EOL>dynamic_shape=initial_state_prior.event_shape_tensor(),<EOL>static_target_shape=[latent_size_],<EOL>dynamic_target_shape=[self.latent_size]),<EOL>_check_equal_shape(<EOL>name="<STR_LIT>",<EOL>static_shape=transition_noise.event_shape,<EOL>dynamic_shape=transition_noise.event_shape_tensor(),<EOL>static_target_shape=[latent_size_],<EOL>dynamic_target_shape=[self.latent_size]),<EOL>_check_equal_shape(<EOL>name="<STR_LIT>",<EOL>static_shape=observation_noise.event_shape,<EOL>dynamic_shape=observation_noise.event_shape_tensor(),<EOL>static_target_shape=[observation_size_],<EOL>dynamic_target_shape=[self.observation_size])]<EOL>self.runtime_assertions = [op for op in runtime_assertions<EOL>if op is not None]<EOL>_, _ = self._batch_shape(), self._batch_shape_tensor()<EOL><DEDENT>super(LinearGaussianStateSpaceModel, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[],<EOL>name=name)<EOL><DEDENT>

Initialize a `LinearGaussianStateSpaceModel.

        Args:
          num_timesteps: Integer `Tensor` total number of timesteps.
          transition_matrix: A transition operator, represented by a Tensor or
            LinearOperator of shape `[latent_size, latent_size]`, or by a
            callable taking as argument a scalar integer Tensor `t` and
            returning a Tensor or LinearOperator representing the transition
            operator from latent state at time `t` to time `t + 1`.
          transition_noise: An instance of
            `tfd.MultivariateNormalLinearOperator` with event shape
            `[latent_size]`, representing the mean and covariance of the
            transition noise model, or a callable taking as argument a
            scalar integer Tensor `t` and returning such a distribution
            representing the noise in the transition from time `t` to time `t + 1`.
          observation_matrix: An observation operator, represented by a Tensor
            or LinearOperator of shape `[observation_size, latent_size]`,
            or by a callable taking as argument a scalar integer Tensor
            `t` and returning a timestep-specific Tensor or
            LinearOperator.
          observation_noise: An instance of
            `tfd.MultivariateNormalLinearOperator` with event shape
            `[observation_size]`, representing the mean and covariance of
            the observation noise model, or a callable taking as argument
            a scalar integer Tensor `t` and returning a timestep-specific
            noise model.
          initial_state_prior: An instance of `MultivariateNormalLinearOperator`
            representing the prior distribution on latent states; must
            have event shape `[latent_size]`.
          initial_step: optional `int` specifying the time of the first
            modeled timestep.  This is added as an offset when passing
            timesteps `t` to (optional) callables specifying
            timestep-specific transition and observation models.
          validate_args: Python `bool`, default `False`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
          allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
          name: The name to give Ops created by the initializer.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>filtered_means = tf.convert_to_tensor(<EOL>value=filtered_means, name="<STR_LIT>")<EOL>filtered_covs = tf.convert_to_tensor(<EOL>value=filtered_covs, name="<STR_LIT>")<EOL>predicted_means = tf.convert_to_tensor(<EOL>value=predicted_means, name="<STR_LIT>")<EOL>predicted_covs = tf.convert_to_tensor(<EOL>value=predicted_covs, name="<STR_LIT>")<EOL>filtered_means = distribution_util.move_dimension(filtered_means, -<NUM_LIT:2>, <NUM_LIT:0>)<EOL>filtered_covs = distribution_util.move_dimension(filtered_covs, -<NUM_LIT:3>, <NUM_LIT:0>)<EOL>predicted_means = distribution_util.move_dimension(predicted_means, -<NUM_LIT:2>, <NUM_LIT:0>)<EOL>predicted_covs = distribution_util.move_dimension(predicted_covs, -<NUM_LIT:3>, <NUM_LIT:0>)<EOL>filtered_means = filtered_means[..., tf.newaxis]<EOL>predicted_means = predicted_means[..., tf.newaxis]<EOL>initial_backward_mean = predicted_means[-<NUM_LIT:1>, ...]<EOL>initial_backward_cov = predicted_covs[-<NUM_LIT:1>, ...]<EOL>num_timesteps = tf.shape(input=filtered_means)[<NUM_LIT:0>]<EOL>initial_state = BackwardPassState(<EOL>backward_mean=initial_backward_mean,<EOL>backward_cov=initial_backward_cov,<EOL>timestep=self.initial_step + num_timesteps - <NUM_LIT:1>)<EOL>update_step_fn = build_backward_pass_step(<EOL>self.get_transition_matrix_for_timestep)<EOL>posterior_states = tf.scan(update_step_fn,<EOL>elems=(filtered_means,<EOL>filtered_covs,<EOL>predicted_means,<EOL>predicted_covs),<EOL>initializer=initial_state,<EOL>reverse=True)<EOL>posterior_means = distribution_util.move_dimension(<EOL>posterior_states.backward_mean[..., <NUM_LIT:0>], <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>posterior_covs = distribution_util.move_dimension(<EOL>posterior_states.backward_cov, <NUM_LIT:0>, -<NUM_LIT:3>)<EOL>return (posterior_means, posterior_covs)<EOL><DEDENT>

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.


def _joint_sample_n(self, n, seed=None):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>stream = seed_stream.SeedStream(<EOL>seed, salt="<STR_LIT>")<EOL>sample_and_batch_shape = distribution_util.prefer_static_value(<EOL>tf.concat([[n], self.batch_shape_tensor()],<EOL>axis=<NUM_LIT:0>))<EOL>with tf.control_dependencies(self.runtime_assertions):<EOL><INDENT>initial_latent = self.initial_state_prior.sample(<EOL>sample_shape=_augment_sample_shape(<EOL>self.initial_state_prior,<EOL>sample_and_batch_shape,<EOL>self.validate_args),<EOL>seed=stream())<EOL>initial_latent = initial_latent[..., tf.newaxis]<EOL><DEDENT>initial_observation_matrix = (<EOL>self.get_observation_matrix_for_timestep(self.initial_step))<EOL>initial_observation_noise = (<EOL>self.get_observation_noise_for_timestep(self.initial_step))<EOL>initial_observation_pred = initial_observation_matrix.matmul(<EOL>initial_latent)<EOL>initial_observation = (initial_observation_pred +<EOL>initial_observation_noise.sample(<EOL>sample_shape=_augment_sample_shape(<EOL>initial_observation_noise,<EOL>sample_and_batch_shape,<EOL>self.validate_args),<EOL>seed=stream())[..., tf.newaxis])<EOL>sample_step = build_kalman_sample_step(<EOL>self.get_transition_matrix_for_timestep,<EOL>self.get_transition_noise_for_timestep,<EOL>self.get_observation_matrix_for_timestep,<EOL>self.get_observation_noise_for_timestep,<EOL>full_sample_and_batch_shape=sample_and_batch_shape,<EOL>stream=stream,<EOL>validate_args=self.validate_args)<EOL>(latents, observations) = tf.scan(<EOL>sample_step,<EOL>elems=tf.range(self.initial_step+<NUM_LIT:1>, self.final_step),<EOL>initializer=(initial_latent, initial_observation))<EOL>latents = tf.concat([initial_latent[tf.newaxis, ...],<EOL>latents], axis=<NUM_LIT:0>)<EOL>observations = tf.concat([initial_observation[tf.newaxis, ...],<EOL>observations], axis=<NUM_LIT:0>)<EOL>latents = tf.squeeze(latents, -<NUM_LIT:1>)<EOL>latents = distribution_util.move_dimension(latents, <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>observations = tf.squeeze(observations, -<NUM_LIT:1>)<EOL>observations = distribution_util.move_dimension(observations, <NUM_LIT:0>, -<NUM_LIT:2>)<EOL><DEDENT>return latents, observations<EOL>

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


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>if mask is not None:<EOL><INDENT>mask = tf.convert_to_tensor(value=mask, name="<STR_LIT>", dtype_hint=tf.bool)<EOL><DEDENT>check_x_shape_op = _check_equal_shape(<EOL>"<STR_LIT:x>", x.shape[-<NUM_LIT:2>:], tf.shape(input=x)[-<NUM_LIT:2>:],<EOL>self.event_shape, self.event_shape_tensor())<EOL>check_mask_dims_op = None<EOL>check_mask_shape_op = None<EOL>if mask is not None:<EOL><INDENT>if (tensorshape_util.rank(mask.shape) is None or<EOL>tensorshape_util.rank(x.shape) is None):<EOL><INDENT>check_mask_dims_op = assert_util.assert_greater_equal(<EOL>tf.rank(x),<EOL>tf.rank(mask),<EOL>message=("<STR_LIT>"))<EOL><DEDENT>elif tensorshape_util.rank(mask.shape) > tensorshape_util.rank(x.shape):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>tensorshape_util.rank(mask.shape),<EOL>tensorshape_util.rank(x.shape)))<EOL><DEDENT>check_mask_shape_op = _check_equal_shape(<EOL>"<STR_LIT>", mask.shape[-<NUM_LIT:1>:], tf.shape(input=mask)[-<NUM_LIT:1>:],<EOL>self.event_shape[-<NUM_LIT:2>:-<NUM_LIT:1>], self.event_shape_tensor()[-<NUM_LIT:2>:-<NUM_LIT:1>])<EOL><DEDENT>if self.validate_args:<EOL><INDENT>runtime_assertions = self.runtime_assertions<EOL>if check_x_shape_op is not None:<EOL><INDENT>runtime_assertions += [check_x_shape_op]<EOL><DEDENT>if check_mask_shape_op is not None:<EOL><INDENT>runtime_assertions += [check_mask_shape_op]<EOL><DEDENT>if check_mask_dims_op is not None:<EOL><INDENT>runtime_assertions += [check_mask_dims_op]<EOL><DEDENT>with tf.control_dependencies(runtime_assertions):<EOL><INDENT>x = tf.identity(x)<EOL><DEDENT><DEDENT>if tensorshape_util.is_fully_defined(<EOL>self.batch_shape) and tensorshape_util.is_fully_defined(x.shape):<EOL><INDENT>sample_and_batch_shape = tf.broadcast_static_shape(<EOL>x.shape[:-<NUM_LIT:2>], self.batch_shape)<EOL><DEDENT>else:<EOL><INDENT>sample_and_batch_shape = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=x)[:-<NUM_LIT:2>], self.batch_shape_tensor())<EOL><DEDENT>if mask is None:<EOL><INDENT>mask_sample_and_batch_shape = self.batch_shape_tensor()<EOL><DEDENT>else:<EOL><INDENT>if (tensorshape_util.is_fully_defined(self.batch_shape) and<EOL>tensorshape_util.is_fully_defined(mask.shape)):<EOL><INDENT>mask_sample_and_batch_shape = tf.broadcast_static_shape(<EOL>mask.shape[:-<NUM_LIT:1>], self.batch_shape)<EOL><DEDENT>else:<EOL><INDENT>mask_sample_and_batch_shape = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=mask)[:-<NUM_LIT:1>], self.batch_shape_tensor())<EOL><DEDENT><DEDENT>x = distribution_util.move_dimension(x, -<NUM_LIT:2>, <NUM_LIT:0>)<EOL>if mask is not None:<EOL><INDENT>mask = distribution_util.move_dimension(mask, -<NUM_LIT:1>, <NUM_LIT:0>)<EOL><DEDENT>x = x[..., tf.newaxis]<EOL>if mask is not None:<EOL><INDENT>mask = mask[..., tf.newaxis, tf.newaxis]<EOL><DEDENT>prior_mean = _broadcast_to_shape(<EOL>self.initial_state_prior.mean()[..., tf.newaxis],<EOL>tf.concat([sample_and_batch_shape,<EOL>[self.latent_size, <NUM_LIT:1>]], axis=<NUM_LIT:0>))<EOL>prior_cov = _broadcast_to_shape(<EOL>self.initial_state_prior.covariance(),<EOL>tf.concat([mask_sample_and_batch_shape,<EOL>[self.latent_size, self.latent_size]], axis=<NUM_LIT:0>))<EOL>initial_observation_matrix = (<EOL>self.get_observation_matrix_for_timestep(self.initial_step))<EOL>initial_observation_noise = (<EOL>self.get_observation_noise_for_timestep(self.initial_step))<EOL>initial_observation_mean = _propagate_mean(prior_mean,<EOL>initial_observation_matrix,<EOL>initial_observation_noise)<EOL>initial_observation_cov = _propagate_cov(prior_cov,<EOL>initial_observation_matrix,<EOL>initial_observation_noise)<EOL>initial_state = KalmanFilterState(<EOL>predicted_mean=prior_mean,<EOL>predicted_cov=prior_cov,<EOL>filtered_mean=prior_mean,  <EOL>filtered_cov=prior_cov,  <EOL>observation_mean=initial_observation_mean,<EOL>observation_cov=initial_observation_cov,<EOL>log_marginal_likelihood=tf.zeros(<EOL>shape=sample_and_batch_shape, dtype=self.dtype),<EOL>timestep=tf.convert_to_tensor(<EOL>value=self.initial_step, dtype=tf.int32, name="<STR_LIT>"))<EOL>update_step_fn = build_kalman_filter_step(<EOL>self.get_transition_matrix_for_timestep,<EOL>self.get_transition_noise_for_timestep,<EOL>self.get_observation_matrix_for_timestep,<EOL>self.get_observation_noise_for_timestep)<EOL>filter_states = tf.scan(update_step_fn,<EOL>elems=x if mask is None else (x, mask),<EOL>initializer=initial_state)<EOL>log_likelihoods = distribution_util.move_dimension(<EOL>filter_states.log_marginal_likelihood, <NUM_LIT:0>, -<NUM_LIT:1>)<EOL>filtered_means = distribution_util.move_dimension(<EOL>filter_states.filtered_mean[..., <NUM_LIT:0>], <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>filtered_covs = distribution_util.move_dimension(<EOL>filter_states.filtered_cov, <NUM_LIT:0>, -<NUM_LIT:3>)<EOL>predicted_means = distribution_util.move_dimension(<EOL>filter_states.predicted_mean[..., <NUM_LIT:0>], <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>predicted_covs = distribution_util.move_dimension(<EOL>filter_states.predicted_cov, <NUM_LIT:0>, -<NUM_LIT:3>)<EOL>observation_means = distribution_util.move_dimension(<EOL>filter_states.observation_mean[..., <NUM_LIT:0>], <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>observation_covs = distribution_util.move_dimension(<EOL>filter_states.observation_cov, <NUM_LIT:0>, -<NUM_LIT:3>)<EOL>return (log_likelihoods,<EOL>filtered_means, filtered_covs,<EOL>predicted_means, predicted_covs,<EOL>observation_means, observation_covs)<EOL><DEDENT>

Run a Kalman filter over a provided sequence of outputs.

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

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

        Returns:
          log_likelihoods: Per-timestep log marginal likelihoods `log
            p(x_t | x_{:t-1})` evaluated at the input `x`, as a `Tensor`
            of shape `sample_shape(x) + batch_shape + [num_timesteps].`
          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
            `sample_shape(mask) + batch_shape + [num_timesteps, latent_size,
            latent_size]`. Note that the covariances depend only on the model and
            the mask, not on the data, so this may have fewer dimensions than
            `filtered_means`.
          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 `sample_shape(mask) + batch_shape +
             [num_timesteps, latent_size, latent_size]`. Note that the covariances
             depend only on the model and the mask, not on the data, so this may
             have fewer dimensions than `predicted_means`.
          observation_means: Means of the per-timestep predictive
             distributions over observations, p(x_{t} | x_{:t-1}), as a
             Tensor of shape `sample_shape(x) + batch_shape +
             [num_timesteps, observation_size]`.
          observation_covs: Covariances of the per-timestep predictive
             distributions over observations, p(x_{t} | x_{:t-1}), as a
             Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
             observation_size, observation_size]`. Note that the covariances depend
             only on the model and the mask, not on the data, so this may have fewer
             dimensions than `observation_means`.


def posterior_marginals(self, x, mask=None):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>(_, filtered_means, filtered_covs,<EOL>predicted_means, predicted_covs, _, _) = self.forward_filter(<EOL>x, mask=mask)<EOL>(smoothed_means, smoothed_covs) = self.backward_smoothing_pass(<EOL>filtered_means, filtered_covs,<EOL>predicted_means, predicted_covs)<EOL>return (smoothed_means, smoothed_covs)<EOL><DEDENT>

Run a Kalman smoother to return posterior mean and cov.

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

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

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

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


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

Compute prior means for all variables via dynamic programming.

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


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

Compute prior covariances for all variables via dynamic programming.

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


def latents_to_observations(self, latent_means, latent_covs):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>pushforward_latents_step = build_pushforward_latents_step(<EOL>self.get_observation_matrix_for_timestep,<EOL>self.get_observation_noise_for_timestep)<EOL>latent_means = distribution_util.move_dimension(<EOL>latent_means, source_idx=-<NUM_LIT:2>, dest_idx=<NUM_LIT:0>)<EOL>latent_means = latent_means[..., tf.newaxis]  <EOL>latent_covs = distribution_util.move_dimension(<EOL>latent_covs, source_idx=-<NUM_LIT:3>, dest_idx=<NUM_LIT:0>)<EOL>(initial_observation_mean,<EOL>initial_observation_cov) = pushforward_latents_step(<EOL>_=None,  <EOL>latent_t_mean_cov=(self.initial_step,<EOL>latent_means[self.initial_step],<EOL>latent_covs[self.initial_step]))<EOL>timesteps = tf.range(self.initial_step,<EOL>self.initial_step + self.num_timesteps)<EOL>observation_means, observation_covs = tf.scan(<EOL>pushforward_latents_step,<EOL>elems=(timesteps, latent_means, latent_covs),<EOL>initializer=(initial_observation_mean, initial_observation_cov),<EOL>parallel_iterations=<NUM_LIT>)<EOL>observation_means = distribution_util.move_dimension(<EOL>observation_means[..., <NUM_LIT:0>], source_idx=<NUM_LIT:0>, dest_idx=-<NUM_LIT:2>)<EOL>observation_covs = distribution_util.move_dimension(<EOL>observation_covs, source_idx=<NUM_LIT:0>, dest_idx=-<NUM_LIT:3>)<EOL>return observation_means, observation_covs<EOL><DEDENT>

Push latent means and covariances forward through the observation model.

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

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


def _compute_non_batch_kl(mu_a, sigma_a, mu_b, sigma_b):

<EOL>sigma_b_inv = np.linalg.inv(sigma_b)<EOL>t = np.trace(sigma_b_inv.dot(sigma_a))<EOL>q = (mu_b - mu_a).dot(sigma_b_inv).dot(mu_b - mu_a)<EOL>k = mu_a.shape[<NUM_LIT:0>]<EOL>l = np.log(np.linalg.det(sigma_b) / np.linalg.det(sigma_a))<EOL>return <NUM_LIT:0.5> * (t + q - k + l)<EOL>

Non-batch KL for N(mu_a, sigma_a), N(mu_b, sigma_b).


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale, low, high], tf.float32)<EOL>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name="<STR_LIT>", dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name="<STR_LIT>", dtype=dtype)<EOL>dtype_util.assert_same_float_dtype([loc, scale, low, high])<EOL>self._broadcast_batch_shape = distribution_util.get_broadcast_shape(<EOL>loc, scale, low, high)<EOL>broadcast_ones = tf.ones(shape=self._broadcast_batch_shape,<EOL>dtype=scale.dtype)<EOL>self._scale = scale * broadcast_ones<EOL>self._loc = loc * broadcast_ones<EOL>self._low = low * broadcast_ones<EOL>self._high = high * broadcast_ones<EOL>with tf.control_dependencies([self._validate()] if validate_args else []):<EOL><INDENT>self._loc = tf.identity(self._loc)<EOL><DEDENT><DEDENT>super(TruncatedNormal, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[loc, scale, low, high],<EOL>name=name)<EOL>

Construct TruncatedNormal.

        All parameters of the distribution will be broadcast to the same shape,
        so the resulting distribution will have a batch_shape of the broadcast
        shape of all parameters.

        Args:
          loc: Floating point tensor; the mean of the normal distribution(s) (
            note that the mean of the resulting distribution will be different
            since it is modified by the bounds).
          scale: Floating point tensor; the std deviation of the normal
            distribution(s).
          low: `float` `Tensor` representing lower bound of the distribution's
            support. Must be such that `low < high`.
          high: `float` `Tensor` representing upper bound of the distribution's
            support. Must be such that `low < high`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked at run-time.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


return self._scale<EOL>


<EOL>ones = tf.ones_like(n - k)<EOL>k_eq_n = tf.equal(k, n)<EOL>safe_dn = tf.where(k_eq_n, ones, n - k)<EOL>dk = tf.math.betainc(a=safe_dn, b=k + <NUM_LIT:1>, x=<NUM_LIT:1> - p)<EOL>return tf.where(k_eq_n, ones, dk)<EOL>

The binomial cumulative distribution function.

    Args:
      k: floating point `Tensor`.
      n: floating point `Tensor`.
      p: floating point `Tensor`.

    Returns:
      `sum_{j=0}^k p^j (1 - p)^(n - j)`.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([total_count, logits, probs], tf.float32)<EOL>self._total_count = self._maybe_assert_valid_total_count(<EOL>tf.convert_to_tensor(<EOL>value=total_count, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits=logits,<EOL>probs=probs,<EOL>validate_args=validate_args,<EOL>name=name,<EOL>dtype=dtype)<EOL><DEDENT>super(Binomial, self).__init__(<EOL>dtype=self._probs.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._total_count, self._logits, self._probs],<EOL>name=name)<EOL>

Initialize a batch of Binomial distributions.

        Args:
          total_count: Non-negative floating point tensor with shape broadcastable
            to `[N1,..., Nm]` with `m >= 0` and the same dtype as `probs` or
            `logits`. Defines this as a batch of `N1 x ...  x Nm` different Binomial
            distributions. Its components should be equal to integer values.
          logits: Floating point tensor representing the log-odds of a
            positive event with shape broadcastable to `[N1,..., Nm]` `m >= 0`, and
            the same dtype as `total_count`. Each entry represents logits for the
            probability of success for independent Binomial distributions. Only one
            of `logits` or `probs` should be passed in.
          probs: Positive floating point tensor with shape broadcastable to
            `[N1,..., Nm]` `m >= 0`, `probs in [0, 1]`. Each entry represents the
            probability of success for independent Binomial distributions. Only one
            of `logits` or `probs` should be passed in.
          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.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


return self._total_count<EOL>


return self._logits<EOL>


return self._probs<EOL>


if not self.validate_args:<EOL><INDENT>return counts<EOL><DEDENT>counts = distribution_util.embed_check_nonnegative_integer_form(counts)<EOL>return distribution_util.with_dependencies([<EOL>assert_util.assert_less_equal(<EOL>counts,<EOL>self.total_count,<EOL>message="<STR_LIT>"),<EOL>], counts)<EOL>

Check counts for proper shape, values, then return tensor version.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>df = tf.convert_to_tensor(<EOL>value=df,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([df], preferred_dtype=tf.float32))<EOL>with tf.control_dependencies([<EOL>assert_util.assert_positive(df),<EOL>] if validate_args else []):<EOL><INDENT>self._df = tf.identity(df, name="<STR_LIT>")<EOL><DEDENT>super(Chi2, self).__init__(<EOL>concentration=<NUM_LIT:0.5> * self._df,<EOL>rate=<NUM_LIT:0.5>,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL><DEDENT>self._parameters = parameters<EOL>

Construct Chi2 distributions with parameter `df`.

        Args:
          df: Floating point tensor, the degrees of freedom of the
            distribution(s). `df` must contain only positive values.
          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.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>return tf.math.log1p(-tf.exp(small - big)) + big<EOL><DEDENT>

Stable evaluation of `Log[exp{big} - exp{small}]`.

    To work correctly, we should have the pointwise relation:  `small <= big`.

    Args:
      big: Floating-point `Tensor`
      small: Floating-point `Tensor` with same `dtype` as `big` and broadcastable
        shape.

    Returns:
      `Tensor` of same `dtype` of `big` and broadcast shape.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._dist = distribution<EOL>if low is not None:<EOL><INDENT>low = tf.convert_to_tensor(<EOL>value=low, name="<STR_LIT>", dtype=distribution.dtype)<EOL><DEDENT>if high is not None:<EOL><INDENT>high = tf.convert_to_tensor(<EOL>value=high, name="<STR_LIT>", dtype=distribution.dtype)<EOL><DEDENT>dtype_util.assert_same_float_dtype(<EOL>tensors=[self.distribution, low, high])<EOL>graph_parents = self._dist._graph_parents  <EOL>checks = []<EOL>if validate_args and low is not None and high is not None:<EOL><INDENT>message = "<STR_LIT>"<EOL>checks.append(assert_util.assert_less(low, high, message=message))<EOL><DEDENT>self._validate_args = validate_args  <EOL>with tf.control_dependencies(checks if validate_args else []):<EOL><INDENT>if low is not None:<EOL><INDENT>self._low = self._check_integer(low)<EOL>graph_parents += [self._low]<EOL><DEDENT>else:<EOL><INDENT>self._low = None<EOL><DEDENT>if high is not None:<EOL><INDENT>self._high = self._check_integer(high)<EOL>graph_parents += [self._high]<EOL><DEDENT>else:<EOL><INDENT>self._high = None<EOL><DEDENT><DEDENT><DEDENT>super(QuantizedDistribution, self).__init__(<EOL>dtype=self._dist.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=self._dist.allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=graph_parents,<EOL>name=name)<EOL>

Construct a Quantized Distribution representing `Y = ceiling(X)`.

        Some properties are inherited from the distribution defining `X`. Example:
        `allow_nan_stats` is determined for this `QuantizedDistribution` by reading
        the `distribution`.

        Args:
          distribution:  The base distribution class to transform. Typically an
            instance of `Distribution`.
          low: `Tensor` with same `dtype` as this distribution and shape
            able to be added to samples. Should be a whole number. Default `None`.
            If provided, base distribution's `prob` should be defined at
            `low`.
          high: `Tensor` with same `dtype` as this distribution and shape
            able to be added to samples. Should be a whole number. Default `None`.
            If provided, base distribution's `prob` should be defined at
            `high - 1`.
            `high` must be strictly greater than `low`.
          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.

        Raises:
          TypeError: If `dist_cls` is not a subclass of
              `Distribution` or continuous.
          NotImplementedError:  If the base distribution does not implement `cdf`.


return self._dist<EOL>


return self._low<EOL>


return self._high<EOL>

Highest value that quantization returns.


<EOL>logsf_y = self.log_survival_function(y)<EOL>logsf_y_minus_1 = self.log_survival_function(y - <NUM_LIT:1>)<EOL>logcdf_y = self.log_cdf(y)<EOL>logcdf_y_minus_1 = self.log_cdf(y - <NUM_LIT:1>)<EOL>big = tf.where(logsf_y < logcdf_y, logsf_y_minus_1, logcdf_y)<EOL>small = tf.where(logsf_y < logcdf_y, logsf_y, logcdf_y_minus_1)<EOL>return _logsum_expbig_minus_expsmall(big, small)<EOL>

Compute log_prob(y) using log survival_function and cdf together.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>if (rate is None) == (log_rate is None):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>elif log_rate is None:<EOL><INDENT>rate = tf.convert_to_tensor(<EOL>value=rate,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([rate], preferred_dtype=tf.float32))<EOL>if not dtype_util.is_floating(rate.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(rate.dtype)))<EOL><DEDENT>with tf.control_dependencies(<EOL>[assert_util.assert_positive(rate)] if validate_args else []):<EOL><INDENT>self._rate = tf.identity(rate, name="<STR_LIT>")<EOL>self._log_rate = tf.math.log(rate, name="<STR_LIT>")<EOL><DEDENT><DEDENT>else:<EOL><INDENT>log_rate = tf.convert_to_tensor(<EOL>value=log_rate,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([log_rate], tf.float32))<EOL>if not dtype_util.is_floating(log_rate.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(log_rate.dtype)))<EOL><DEDENT>self._rate = tf.exp(log_rate, name="<STR_LIT>")<EOL>self._log_rate = tf.convert_to_tensor(value=log_rate, name="<STR_LIT>")<EOL><DEDENT><DEDENT>self._interpolate_nondiscrete = interpolate_nondiscrete<EOL>super(Poisson, self).__init__(<EOL>dtype=self._rate.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._rate],<EOL>name=name)<EOL>

Initialize a batch of Poisson distributions.

        Args:
          rate: Floating point tensor, the rate parameter. `rate` must be positive.
            Must specify exactly one of `rate` and `log_rate`.
          log_rate: Floating point tensor, the log of the rate parameter.
            Must specify exactly one of `rate` and `log_rate`.
          interpolate_nondiscrete: Python `bool`. When `False`,
            `log_prob` returns `-inf` (and `prob` returns `0`) for non-integer
            inputs. When `True`, `log_prob` evaluates the continuous function
            `k * log_rate - lgamma(k+1) - rate`, which matches the Poisson pmf
            at integer arguments `k` (note that this function is not itself
            a normalized probability log-density).
            Default value: `True`.
          validate_args: Python `bool`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          allow_nan_stats: Python `bool`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
            Default value: `True`.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: if none or both of `rate`, `log_rate` are specified.
          TypeError: if `rate` is not a float-type.
          TypeError: if `log_rate` is not a float-type.


return self._rate<EOL>


return self._log_rate<EOL>


return self._interpolate_nondiscrete<EOL>

Interpolate (log) probs on non-integer inputs.


self._seed = seed.original_seed if isinstance(seed, SeedStream) else seed<EOL>self._salt = salt<EOL>self._counter = <NUM_LIT:0><EOL>

Initializes a `SeedStream`.

        Args:
          seed: Any Python object convertible to string, supplying the
            initial entropy.  If `None`, operations seeded with seeds
            drawn from this `SeedStream` will follow TensorFlow semantics
            for not being seeded.
          salt: Any Python object convertible to string, supplying
            auxiliary entropy.  Must be unique across the Distributions
            and TensorFlow Probability code base.  See class docstring for
            rationale.


self._counter += <NUM_LIT:1><EOL>if self._seed is None:<EOL><INDENT>return None<EOL><DEDENT>composite = str((self._seed, self._counter, self._salt)).encode('<STR_LIT:utf-8>')<EOL>return int(hashlib.sha512(composite).hexdigest(), <NUM_LIT:16>)<EOL>

Returns a fresh integer usable as a seed in downstream operations.

        If this `SeedStream` was initialized with `seed=None`, returns
        `None`.  This has the effect that downstream operations (both
        `SeedStream`s and primitive TensorFlow ops) will behave as though
        they were unseeded.

        The returned integer is non-negative, and uniformly distributed in
        the half-open interval `[0, 2**512)`.  This is consistent with
        TensorFlow, as TensorFlow operations internally use the residue of
        the given seed modulo `2**31 - 1` (see
        `tensorflow/python/framework/random_seed.py`).

        Returns:
          seed: A fresh integer usable as a seed in downstream operations,
            or `None`.


def __init__(self, model, sample_dtype=None, validate_args=False, name=None):

parameters = dict(locals())<EOL>with tf.name_scope(name or '<STR_LIT>') as name:<EOL><INDENT>self._sample_dtype = sample_dtype<EOL>self._model = model<EOL>self._most_recently_built_distributions = None<EOL>super(JointDistributionCoroutine, self).__init__(<EOL>dtype=sample_dtype,<EOL>reparameterization_type=None,  <EOL>validate_args=validate_args,<EOL>allow_nan_stats=False,<EOL>parameters=parameters,<EOL>graph_parents=[],<EOL>name=name)<EOL><DEDENT>

Construct the `JointDistributionCoroutine` distribution.

        Args:
          model: A generator that yields a sequence of `tfd.Distribution`-like
            instances.
          sample_dtype: Samples from this distribution will be structured like
            `tf.nest.pack_sequence_as(sample_dtype, list_)`. `sample_dtype` is only
            used for `tf.nest.pack_sequence_as` structuring of outputs, never
            casting (which is the responsibility of the component distributions).
            Default value: `None` (i.e., `tuple`).
          validate_args: Python `bool`.  Whether to validate input with asserts.
            If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
            Default value: `False`.
          name: The name for ops managed by the distribution.
            Default value: `None` (i.e., `"JointDistributionCoroutine"`).


def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):

ds = []<EOL>values_out = []<EOL>seed = seed_stream.SeedStream('<STR_LIT>', seed)<EOL>gen = self._model()<EOL>index = <NUM_LIT:0><EOL>d = next(gen)<EOL>try:<EOL><INDENT>while True:<EOL><INDENT>actual_distribution = d.distribution if isinstance(d, self.Root) else d<EOL>ds.append(actual_distribution)<EOL>if (value is not None and len(value) > index and<EOL>value[index] is not None):<EOL><INDENT>seed()<EOL>next_value = value[index]<EOL><DEDENT>else:<EOL><INDENT>next_value = actual_distribution.sample(<EOL>sample_shape=sample_shape if isinstance(d, self.Root) else (),<EOL>seed=seed())<EOL><DEDENT>values_out.append(next_value)<EOL>index += <NUM_LIT:1><EOL>d = gen.send(next_value)<EOL><DEDENT><DEDENT>except StopIteration:<EOL><INDENT>pass<EOL><DEDENT>return ds, values_out<EOL>

Executes `model`, creating both samples and distributions.


def _verifyPdfWithNumpy(self, vmf, atol=<NUM_LIT>):

dim = tf.compat.dimension_value(vmf.event_shape[-<NUM_LIT:1>])<EOL>nsamples = <NUM_LIT:10><EOL>sample_shape = [nsamples] + tensorshape_util.as_list(<EOL>vmf.batch_shape) + [dim]<EOL>uniforms = np.random.randn(*sample_shape)<EOL>uniforms /= np.linalg.norm(uniforms, axis=-<NUM_LIT:1>, keepdims=True)<EOL>uniforms = uniforms.astype(dtype_util.as_numpy_dtype(vmf.dtype))<EOL>vmf_samples = vmf.sample(<EOL>sample_shape=[nsamples], seed=tfp_test_util.test_seed())<EOL>samples = tf.concat([uniforms, vmf_samples], axis=<NUM_LIT:0>)<EOL>samples = tf.debugging.check_numerics(samples, '<STR_LIT>')<EOL>samples = self.evaluate(samples)<EOL>log_prob = vmf.log_prob(samples)<EOL>log_prob = tf.debugging.check_numerics(log_prob, '<STR_LIT>')<EOL>try:<EOL><INDENT>from scipy.special import gammaln  <EOL>from scipy.special import ive  <EOL><DEDENT>except ImportError:<EOL><INDENT>tf.compat.v1.logging.warn('<STR_LIT>')<EOL>return<EOL><DEDENT>conc = self.evaluate(vmf.concentration)<EOL>mean_dir = self.evaluate(vmf.mean_direction)<EOL>log_true_sphere_surface_area = (<EOL>np.log(<NUM_LIT:2>) + (dim / <NUM_LIT:2>) * np.log(np.pi) - gammaln(dim / <NUM_LIT:2>))<EOL>expected = (<EOL>conc * np.sum(samples * mean_dir, axis=-<NUM_LIT:1>) +<EOL>np.where(conc > <NUM_LIT:0>,<EOL>(dim / <NUM_LIT:2> - <NUM_LIT:1>) * np.log(conc) -<EOL>(dim / <NUM_LIT:2>) * np.log(<NUM_LIT:2> * np.pi) -<EOL>np.log(ive(dim / <NUM_LIT:2> - <NUM_LIT:1>, conc)) -<EOL>np.abs(conc),<EOL>-log_true_sphere_surface_area))<EOL>self.assertAllClose(expected, self.evaluate(log_prob),<EOL>atol=atol)<EOL>

Verifies log_prob evaluations with numpy/scipy.

        Both uniform random points and sampled points are evaluated.

        Args:
          vmf: A `tfp.distributions.VonMisesFisher` instance.
          atol: Absolute difference tolerable.


def _verifySampleAndPdfConsistency(self, vmf, rtol=<NUM_LIT>):

dim = tf.compat.dimension_value(vmf.event_shape[-<NUM_LIT:1>])<EOL>nsamples = <NUM_LIT><EOL>samples = vmf.sample(<EOL>sample_shape=[nsamples], seed=tfp_test_util.test_seed())<EOL>samples = tf.debugging.check_numerics(samples, '<STR_LIT>')<EOL>log_prob = vmf.log_prob(samples)<EOL>log_prob = tf.debugging.check_numerics(log_prob, '<STR_LIT>')<EOL>log_importance = -log_prob<EOL>sphere_surface_area_estimate, samples, importance, conc = self.evaluate([<EOL>tf.exp(<EOL>tf.reduce_logsumexp(input_tensor=log_importance, axis=<NUM_LIT:0>) -<EOL>tf.math.log(tf.cast(nsamples, dtype=tf.float32))), samples,<EOL>tf.exp(log_importance), vmf.concentration<EOL>])<EOL>true_sphere_surface_area = <NUM_LIT:2> * (np.pi)**(dim / <NUM_LIT:2>) * self.evaluate(<EOL>tf.exp(-tf.math.lgamma(dim / <NUM_LIT:2>)))<EOL>true_sphere_surface_area += np.zeros_like(sphere_surface_area_estimate)<EOL>self.assertAllClose(<EOL>true_sphere_surface_area[np.where(conc < <NUM_LIT:3>)],<EOL>sphere_surface_area_estimate[np.where(conc < <NUM_LIT:3>)],<EOL>rtol=rtol)<EOL>for survival_rate in <NUM_LIT>, <NUM_LIT>, <NUM_LIT>, <NUM_LIT>:<EOL><INDENT>cdf = (<NUM_LIT:1> - survival_rate)<EOL>mean_dir = self.evaluate(vmf.mean_direction)<EOL>dotprods = np.sum(samples * mean_dir, -<NUM_LIT:1>)<EOL>dotprod_thresh = np.percentile(<EOL>dotprods, <NUM_LIT:100> * survival_rate, axis=<NUM_LIT:0>, keepdims=True)<EOL>dotprod_above_thresh = np.float32(dotprods > dotprod_thresh)<EOL>sphere_cap_surface_area_ests = (<EOL>cdf * (importance * dotprod_above_thresh).sum(<NUM_LIT:0>) /<EOL>dotprod_above_thresh.sum(<NUM_LIT:0>))<EOL>h = (<NUM_LIT:1> - dotprod_thresh)<EOL>self.assertGreaterEqual(h.min(), <NUM_LIT:0>)  <EOL>true_sphere_cap_surface_area = (<EOL><NUM_LIT:0.5> * true_sphere_surface_area *<EOL>self.evaluate(tf.math.betainc((dim - <NUM_LIT:1>) / <NUM_LIT:2>, <NUM_LIT:0.5>, <NUM_LIT:2> * h - h**<NUM_LIT:2>)))<EOL>if dim == <NUM_LIT:3>:  <EOL><INDENT>self.assertAllClose(<NUM_LIT:2> * np.pi * h, true_sphere_cap_surface_area)<EOL><DEDENT>self.assertAllClose(<EOL>true_sphere_cap_surface_area,<EOL>sphere_cap_surface_area_ests +<EOL>np.zeros_like(true_sphere_cap_surface_area),<EOL>rtol=rtol)<EOL><DEDENT>

Verifies samples are consistent with the PDF using importance sampling.

        In particular, we verify an estimate the surface area of the n-dimensional
        hypersphere, and the surface areas of the spherical caps demarcated by
        a handful of survival rates.

        Args:
          vmf: A `VonMisesFisher` distribution instance.
          rtol: Relative difference tolerable.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([scale],<EOL>preferred_dtype=tf.float32)<EOL>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(scale)] if validate_args else []):<EOL><INDENT>self._scale = tf.identity(<EOL>scale, name="<STR_LIT>")<EOL><DEDENT><DEDENT>self._half_cauchy = half_cauchy.HalfCauchy(<EOL>loc=tf.zeros([], dtype=dtype),<EOL>scale=tf.ones([], dtype=dtype),<EOL>allow_nan_stats=True)<EOL>super(Horseshoe, self).__init__(<EOL>dtype=self._scale.dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._scale],<EOL>name=name)<EOL>

Construct a Horseshoe distribution with `scale`.

        Args:
          scale: Floating point tensor; the scales of the distribution(s).
            Must contain only positive values.
          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. Default value: `False` (i.e., do not validate args).
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or more
            of the statistic's batch members are undefined.
            Default value: `True`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: 'Horseshoe'.


return self._scale<EOL>


one = tf.ones([], dtype=x.dtype)<EOL>return tf.math.digamma(x + one) - tf.math.digamma(one)<EOL>

Compute the harmonic number from its analytic continuation.

    Derivation from [here](
    https://en.wikipedia.org/wiki/Digamma_function#Relation_to_harmonic_numbers)
    and [Euler's constant](
    https://en.wikipedia.org/wiki/Euler%E2%80%93Mascheroni_constant).

    Args:
      x: input float.

    Returns:
      z: The analytic continuation of the harmonic number for the input.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([concentration1, concentration0],<EOL>tf.float32)<EOL>concentration1 = tf.convert_to_tensor(<EOL>value=concentration1, name="<STR_LIT>", dtype=dtype)<EOL>concentration0 = tf.convert_to_tensor(<EOL>value=concentration0, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>super(Kumaraswamy, self).__init__(<EOL>distribution=uniform.Uniform(<EOL>low=tf.zeros([], dtype=concentration1.dtype),<EOL>high=tf.ones([], dtype=concentration1.dtype),<EOL>allow_nan_stats=allow_nan_stats),<EOL>bijector=kumaraswamy_bijector.Kumaraswamy(<EOL>concentration1=concentration1,<EOL>concentration0=concentration0,<EOL>validate_args=validate_args),<EOL>batch_shape=distribution_util.get_broadcast_shape(<EOL>concentration1, concentration0),<EOL>parameters=parameters,<EOL>name=name)<EOL>

Initialize a batch of Kumaraswamy distributions.

        Args:
          concentration1: Positive floating-point `Tensor` indicating mean
            number of successes; aka "alpha". Implies `self.dtype` and
            `self.batch_shape`, i.e.,
            `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
          concentration0: Positive floating-point `Tensor` indicating mean
            number of failures; aka "beta". Otherwise has same semantics as
            `concentration1`.
          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.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


return self.bijector.concentration1<EOL>

Concentration parameter associated with a `1` outcome.


return self.bijector.concentration0<EOL>

Concentration parameter associated with a `0` outcome.


total_concentration = self.concentration1 + self.concentration0<EOL>expanded_concentration1 = tf.ones_like(<EOL>total_concentration, dtype=self.dtype) * self.concentration1<EOL>expanded_concentration0 = tf.ones_like(<EOL>total_concentration, dtype=self.dtype) * self.concentration0<EOL>beta_arg0 = <NUM_LIT:1> + n / expanded_concentration1<EOL>beta_arg = tf.stack([beta_arg0, expanded_concentration0], -<NUM_LIT:1>)<EOL>log_moment = tf.math.log(expanded_concentration0) + tf.math.lbeta(beta_arg)<EOL>return tf.exp(log_moment)<EOL>


def _broadcast_cat_event_and_params(event, params, base_dtype):

if dtype_util.is_integer(event.dtype):<EOL><INDENT>pass<EOL><DEDENT>elif dtype_util.is_floating(event.dtype):<EOL><INDENT>event = tf.cast(event, dtype=tf.int32)<EOL><DEDENT>else:<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(base_dtype))<EOL><DEDENT>shape_known_statically = (<EOL>tensorshape_util.rank(params.shape) is not None and<EOL>tensorshape_util.is_fully_defined(params.shape[:-<NUM_LIT:1>]) and<EOL>tensorshape_util.is_fully_defined(event.shape))<EOL>if not shape_known_statically or params.shape[:-<NUM_LIT:1>] != event.shape:<EOL><INDENT>params *= tf.ones_like(event[..., tf.newaxis],<EOL>dtype=params.dtype)<EOL>params_shape = tf.shape(input=params)[:-<NUM_LIT:1>]<EOL>event *= tf.ones(params_shape, dtype=event.dtype)<EOL>if tensorshape_util.rank(params.shape) is not None:<EOL><INDENT>tensorshape_util.set_shape(event, params.shape[:-<NUM_LIT:1>])<EOL><DEDENT><DEDENT>return event, params<EOL>

Broadcasts the event or distribution parameters.


@kullback_leibler.RegisterKL(Categorical, Categorical)<EOL>def _kl_categorical_categorical(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>delta_log_probs1 = (tf.nn.log_softmax(a.logits) -<EOL>tf.nn.log_softmax(b.logits))<EOL>return tf.reduce_sum(<EOL>input_tensor=tf.nn.softmax(a.logits) * delta_log_probs1, axis=-<NUM_LIT:1>)<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b Categorical.

    Args:
      a: instance of a Categorical distribution object.
      b: instance of a Categorical distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_categorical_categorical".

    Returns:
      Batchwise KL(a || b)