text: # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """The Sample distribution class."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import functools
  
  # Dependency imports
  import numpy as np
  
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python import math as tfp_math
  from tensorflow_probability.python.bijectors import bijector as bijector_lib
  from tensorflow_probability.python.distributions import distribution as distribution_lib
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static as ps
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  
  
  def _make_summary_statistic(attr):
    """Factory for implementing summary statistics, eg, mean, stddev, mode."""
    def _fn(self, **kwargs):
      """Implements summary statistic, eg, mean, stddev, mode."""
      sample_shape = ps.reshape(self.sample_shape, shape=[-1])
      x = getattr(self.distribution, attr)(**kwargs)
      shape = ps.concat([
          self.distribution.batch_shape_tensor(),
          ps.ones(ps.rank_from_shape(sample_shape), dtype=sample_shape.dtype),
          self.distribution.event_shape_tensor(),
      ], axis=0)
      x = tf.reshape(x, shape=shape)
      shape = ps.concat([
          self.distribution.batch_shape_tensor(),
          sample_shape,
          self.distribution.event_shape_tensor(),
      ], axis=0)
      return tf.broadcast_to(x, shape)
    return _fn
  
  
  class Sample(distribution_lib.Distribution):
    """Distribution over IID samples of a given shape.
  
    Given random variable `X`, one may make a new random variable by concatenating
    samples.  For example, if `X1` and `X2` are iid `Normal(0, 1)` samples then
    `[X1, X2]` is a bi-variate normal random vector.
  
    #### Mathematical Details
  
    With `p` the probability density/mass of the function being sampled, and
    `n` the samples taken, the density/mass of this distribution is
  
    ```none
    p(x) = prod{ p(x[i]) : i = 0, ..., (n - 1) }
    ```
  
    #### Examples
  
    ```python
    tfd = tfp.distributions
  
    # Example 1: Five scalar draws.
  
    s = tfd.Sample(
          tfd.Normal(loc=0, scale=1),
          sample_shape=5)
    x = s.sample()
    # ==> x.shape: [5]
  
    lp = s.log_prob(x)
    # ==> lp.shape: []
    #     Equivalently: tf.reduce_sum(s.distribution.log_prob(x), axis=0)
    #
    # `Sample.log_prob` computes the per-{sample, batch} `log_prob`s then sums
    # over the `Sample.sample_shape` dimensions. In the above example `log_prob`
    # dims `0` is summed out, since it is the `sample_shape` dimension.
  
    # Example 2: `[5, 4]`-draws of a bivariate Normal.
  
    mvn = tfd.MultivariateNormalDiag(loc=tf.zeros([3, 2]))
    mvn.batch_shape ==> [3]
    mvn.event_shape ==> [2]
  
    s = tfd.Sample(mvn, sample_shape=[5, 4])
    s.batch_shape ==> [3]
    s.event_shape ==> [5, 4, 2]
  
    x = s.sample([6, 1])
    # ==> x.shape: [6, 1, 3, 5, 4, 2]
  
    lp = s.log_prob(x)
    # ==> lp.shape: [6, 1, 3]
    #
    # `s.log_prob` will reduce over the event dims of `mvn`, i.e., dim `5`, then
    # sums over `s.sample_shape` dims `[3, 4]` corresponding to shape (slice)
    # `[5, 4]`.
    ```
  
    """
  
    def __init__(
        self,
        distribution,
        sample_shape=(),
        validate_args=False,
        experimental_use_kahan_sum=False,
        name=None):
      """Construct the `Sample` distribution.
  
      The `event_shape` and `batch_shape` of the `Sample` distribution are
      determined by the args `distribution` and `sample_shape`:
  
      ```
      s = Sample(distribution, sample_shape)
      ==> s.batch_shape: distribution.batch_shape
      ==> s.event_shape: sample_shape + distribution.event_shape
      ```
  
      Args:
        distribution: The base distribution instance to transform. Typically an
          instance of `Distribution`.
        sample_shape: `int` scalar or vector `Tensor` representing the shape of a
          single sample.
        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.
        experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
          summation to aggregate independent underlying log_prob values, which
          improves against the precision of a naive float32 sum. This can be
          noticeable in particular for large dimensions in float32. See CPU caveat
          on `tfp.math.reduce_kahan_sum`.
        name: The name for ops managed by the distribution.
          Default value: `None` (i.e., `'Sample' + distribution.name`).
      """
      parameters = dict(locals())
      self._experimental_use_kahan_sum = experimental_use_kahan_sum
      with tf.name_scope(name or 'Sample' + distribution.name) as name:
        self._distribution = distribution
        self._sample_shape = tensor_util.convert_nonref_to_tensor(
            sample_shape, dtype_hint=tf.int32, name='sample_shape',
            as_shape_tensor=True)
        super(Sample, self).__init__(
            dtype=self._distribution.dtype,
            reparameterization_type=self._distribution.reparameterization_type,
            validate_args=validate_args,
            allow_nan_stats=self._distribution.allow_nan_stats,
            parameters=parameters,
            name=name)
  
    @property
    def distribution(self):
      return self._distribution
  
    @property
    def sample_shape(self):
      return self._sample_shape
  
    @property
    def experimental_is_sharded(self):
      return self.distribution.experimental_is_sharded
  
    def __getitem__(self, slices):
      # Because slicing is parameterization-dependent, we only implement slicing
      # for instances of MSF, not subclasses thereof.
      if type(self).__init__ is not Sample.__init__:
        return super(Sample, self).__getitem__(slices)
  
      # Since this distribution only modifies the event shape, we can simply pass
      # through slicing to the underlying.
      return self.copy(distribution=self.distribution.__getitem__(slices))
  
    def _batch_shape_tensor(self):
      return self.distribution.batch_shape_tensor()
  
    def _batch_shape(self):
      return self.distribution.batch_shape
  
    def _event_shape_tensor(self):
      return ps.concat([
          ps.reshape(self.sample_shape, shape=[-1]),
          self.distribution.event_shape_tensor(),
      ], axis=0)
  
    def _event_shape(self):
      s = tf.get_static_value(self.sample_shape)
      if tensorshape_util.rank(s) == 1:
        sample_shape = tf.TensorShape(s)
      else:
        sample_shape = tensorshape_util.constant_value_as_shape(self.sample_shape)
      if (tensorshape_util.rank(sample_shape) is None or
          tensorshape_util.rank(self.distribution.event_shape) is None):
        return tf.TensorShape(None)
      return tensorshape_util.concatenate(sample_shape,
                                          self.distribution.event_shape)
  
    def _sample_n(self, n, seed, **kwargs):
      sample_shape = ps.reshape(self.sample_shape, shape=[-1])
      fake_sample_ndims = ps.rank_from_shape(sample_shape)
      event_ndims = ps.rank_from_shape(
          self.distribution.event_shape_tensor, self.distribution.event_shape)
      batch_ndims = ps.rank_from_shape(
          self.distribution.batch_shape_tensor, self.distribution.batch_shape)
      perm = ps.concat([
          [0],
          ps.range(1 + fake_sample_ndims,
                   1 + fake_sample_ndims + batch_ndims,
                   dtype=tf.int32),
          ps.range(1, 1 + fake_sample_ndims, dtype=tf.int32),
          ps.range(1 + fake_sample_ndims + batch_ndims,
                   1 + fake_sample_ndims + batch_ndims + event_ndims,
                   dtype=tf.int32),
      ], axis=0)
      x = self.distribution.sample(
          ps.concat([[n], sample_shape], axis=0),
          seed=seed,
          **kwargs)
      return tf.transpose(a=x, perm=perm)
  
    def _sum_fn(self):
      if self._experimental_use_kahan_sum:
        return lambda x, axis: tfp_math.reduce_kahan_sum(x, axis).total
      return tf.math.reduce_sum
  
    def _prepare_for_underlying(self, x):
      batch_ndims = ps.rank_from_shape(
          self.distribution.batch_shape_tensor,
          self.distribution.batch_shape)
      extra_sample_ndims = ps.rank_from_shape(self.sample_shape)
      event_ndims = ps.rank_from_shape(
          self.distribution.event_shape_tensor,
          self.distribution.event_shape)
      ndims = ps.rank(x)
      # (1) Expand x's dims.
      d = ndims - batch_ndims - extra_sample_ndims - event_ndims
      x = tf.reshape(
          x,
          shape=ps.pad(
              ps.shape(x),
              paddings=[[ps.maximum(0, -d), 0]],
              constant_values=1))
      ndims = ps.rank(x)
      sample_ndims = ps.maximum(0, d)
      # (2) Transpose x's dims.
      sample_dims = ps.range(0, sample_ndims)
      batch_dims = ps.range(sample_ndims, sample_ndims + batch_ndims)
      extra_sample_dims = ps.range(
          sample_ndims + batch_ndims,
          sample_ndims + batch_ndims + extra_sample_ndims)
      event_dims = ps.range(
          sample_ndims + batch_ndims + extra_sample_ndims,
          ndims)
      perm = ps.concat(
          [sample_dims, extra_sample_dims, batch_dims, event_dims], axis=0)
      x = tf.transpose(x, perm=perm)
      return x, (sample_ndims, extra_sample_ndims, batch_ndims)
  
    def _finish_log_prob(self, lp, aux):
      (sample_ndims, extra_sample_ndims, batch_ndims) = aux
      # (1) Ensure lp is fully broadcast in the sample dims, i.e. ensure lp has
      #     full sample shape in the sample axes, before we reduce.
      bcast_lp_shape = ps.broadcast_shape(
          ps.shape(lp),
          ps.concat([ps.ones([sample_ndims], tf.int32),
                     ps.reshape(self.sample_shape, shape=[-1]),
                     ps.ones([batch_ndims], tf.int32)], axis=0))
      lp = tf.broadcast_to(lp, bcast_lp_shape)
      # (2) Make the final reduction.
      axis = ps.range(sample_ndims, sample_ndims + extra_sample_ndims)
      return self._sum_fn()(lp, axis=axis)
  
    def _log_prob(self, x, **kwargs):
      x, aux = self._prepare_for_underlying(x)
      return self._finish_log_prob(
          self.distribution.log_prob(x, **kwargs),
          aux)
  
    def _entropy(self, **kwargs):
      h = self.distribution.entropy(**kwargs)
      n = ps.reduce_prod(self.sample_shape)
      return tf.cast(x=n, dtype=h.dtype) * h
  
    _mean = _make_summary_statistic('mean')
    _stddev = _make_summary_statistic('stddev')
    _variance = _make_summary_statistic('variance')
    _mode = _make_summary_statistic('mode')
  
    def _default_event_space_bijector(self):
      # TODO(b/170405182): In scenarios where we can statically prove that it has
      #   no batch part, avoid the transposes by directly using
      #   `self.distribution.experimental_default_event_space_bijector()`.
      bijector = self.distribution.experimental_default_event_space_bijector()
      if bijector is None:
        return None
      return _DefaultSampleBijector(self.distribution, self.sample_shape,
                                    self._sum_fn(), bijector=bijector)
  
    def _parameter_control_dependencies(self, is_init):
      assertions = []
      sample_shape = None  # Memoize concretization.
  
      # Check valid shape.
      ndims_ = tensorshape_util.rank(self.sample_shape.shape)
      if is_init != (ndims_ is None):
        msg = 'Argument `sample_shape` must be either a scalar or a vector.'
        if ndims_ is not None:
          if ndims_ > 1:
            raise ValueError(msg)
        elif self.validate_args:
          if sample_shape is None:
            sample_shape = tf.convert_to_tensor(self.sample_shape)
          assertions.append(assert_util.assert_less(
              tf.rank(sample_shape), 2, message=msg))
  
      # Check valid dtype.
      if is_init:  # No xor check because `dtype` cannot change.
        dtype_ = self.sample_shape.dtype
        if dtype_ is None:
          if sample_shape is None:
            sample_shape = tf.convert_to_tensor(self.sample_shape)
          dtype_ = sample_shape.dtype
        if dtype_util.base_dtype(dtype_) not in {tf.int32, tf.int64}:
          raise TypeError('Argument `sample_shape` must be integer type; '
                          'saw {}.'.format(dtype_util.name(dtype_)))
  
      # Check valid "value".
      if is_init != tensor_util.is_ref(self.sample_shape):
        sample_shape_ = tf.get_static_value(self.sample_shape)
        msg = 'Argument `sample_shape` must have non-negative values.'
        if sample_shape_ is not None:
          if np.any(np.array(sample_shape_) < 0):
            raise ValueError('{} Saw: {}'.format(msg, sample_shape_))
        elif self.validate_args:
          if sample_shape is None:
            sample_shape = tf.convert_to_tensor(self.sample_shape)
          assertions.append(assert_util.assert_greater(
              sample_shape, -1, message=msg))
  
      return assertions
  
    _composite_tensor_nonshape_params = ('distribution',)
  
    _composite_tensor_shape_params = ('sample_shape',)
  
  
  class _DefaultSampleBijector(bijector_lib.Bijector):
    """Since tfd.Sample uses transposes, it requires a custom event bijector."""
  
    def __init__(self, distribution, sample_shape, sum_fn, bijector=None):
      parameters = dict(locals())
      self.distribution = distribution
      if bijector is None:
        bijector = distribution.experimental_default_event_space_bijector()
      self.bijector = bijector
      self.sample_shape = sample_shape
      self._sum_fn = sum_fn
      sample_ndims = ps.rank_from_shape(self.sample_shape)
      super(_DefaultSampleBijector, self).__init__(
          forward_min_event_ndims=(
              self.bijector.forward_min_event_ndims + sample_ndims),
          inverse_min_event_ndims=(
              self.bijector.inverse_min_event_ndims + sample_ndims),
          parameters=parameters)
  
    def _forward_event_shape(self, shape):
      return self.bijector.forward_event_shape(shape)
  
    def _forward_event_shape_tensor(self, shape):
      return self.bijector.forward_event_shape_tensor(shape)
  
    def _inverse_event_shape(self, shape):
      return self.bijector.inverse_event_shape(shape)
  
    def _inverse_event_shape_tensor(self, shape):
      return self.bijector.inverse_event_shape_tensor(shape)
  
    def _transpose_around_bijector_fn(self,
                                      bijector_fn,
                                      arg,
                                      src_event_ndims,
                                      dest_event_ndims=None,
                                      fn_reduces_event=False,
                                      **kwargs):
      # This function moves the axes corresponding to `self.sample_shape` to the
      # left of the batch shape, then applies `bijector_fn`, then moves the axes
      # corresponding to `self.sample_shape` back to the event part of the shape.
      #
      # `src_event_ndims` and `dest_event_ndims` indicate the expected event rank
      # (omitting `self.sample_shape`) before and after applying `bijector_fn`.
      #
      # This function arose because forward and inverse ended up being quite
      # similar. It was then only a small generalization to also support {F/I}LDJ.
      batch_ndims = ps.rank_from_shape(self.distribution.batch_shape_tensor,
                                       self.distribution.batch_shape)
      extra_sample_ndims = ps.rank_from_shape(self.sample_shape)
      arg_ndims = ps.rank(arg)
      # (1) Expand arg's dims.
      d = arg_ndims - batch_ndims - extra_sample_ndims - src_event_ndims
      arg = tf.reshape(
          arg,
          shape=ps.pad(
              ps.shape(arg),
              paddings=[[ps.maximum(0, -d), 0]],
              constant_values=1))
      arg_ndims = ps.rank(arg)
      sample_ndims = ps.maximum(0, d)
      # (2) Transpose arg's dims.
      sample_dims = ps.range(0, sample_ndims)
      batch_dims = ps.range(sample_ndims, sample_ndims + batch_ndims)
      extra_sample_dims = ps.range(
          sample_ndims + batch_ndims,
          sample_ndims + batch_ndims + extra_sample_ndims)
      event_dims = ps.range(
          sample_ndims + batch_ndims + extra_sample_ndims,
          arg_ndims)
      perm = ps.concat(
          [sample_dims, extra_sample_dims, batch_dims, event_dims], axis=0)
      arg = tf.transpose(arg, perm=perm)
      # (3) Apply underlying bijector.
      result = bijector_fn(arg, **kwargs)
      # (4) Transpose sample_shape from the sample to the event shape.
      result_ndims = ps.rank(result)
      if fn_reduces_event:
        dest_event_ndims = 0
      d = result_ndims - batch_ndims - extra_sample_ndims - dest_event_ndims
      if fn_reduces_event:
        # In some cases, fn may reduce event too far, i.e. ildj may return a
        # scalar `0.`, which won't work with the transpose we do below.
        result = tf.reshape(
            result,
            shape=ps.pad(
                ps.shape(result),
                paddings=[[ps.maximum(0, -d), 0]],
                constant_values=1))
        result_ndims = ps.rank(result)
      sample_ndims = ps.maximum(0, d)
      sample_dims = ps.range(0, sample_ndims)
      extra_sample_dims = ps.range(sample_ndims,
                                   sample_ndims + extra_sample_ndims)
      batch_dims = ps.range(sample_ndims + extra_sample_ndims,
                            sample_ndims + extra_sample_ndims + batch_ndims)
      event_dims = ps.range(sample_ndims + extra_sample_ndims + batch_ndims,
                            result_ndims)
      perm = ps.concat(
          [sample_dims, batch_dims, extra_sample_dims, event_dims], axis=0)
      return tf.transpose(result, perm=perm)
  
    def _forward(self, x, **kwargs):
      dist = self.distribution
      event_ndims = ps.rank_from_shape(dist.event_shape_tensor, dist.event_shape)
      bij = self.bijector
      pullback_event_ndims = ps.rank_from_shape(
          lambda: bij.inverse_event_shape_tensor(dist.event_shape_tensor()),
          bij.inverse_event_shape(dist.event_shape))
      return self._transpose_around_bijector_fn(
          bij.forward, arg=x,
          src_event_ndims=pullback_event_ndims, dest_event_ndims=event_ndims)
  
    def _inverse(self, y, **kwargs):
      dist = self.distribution
      event_ndims = ps.rank_from_shape(dist.event_shape_tensor, dist.event_shape)
      bij = self.bijector
      pullback_event_ndims = ps.rank_from_shape(
          lambda: bij.inverse_event_shape_tensor(dist.event_shape_tensor()),
          bij.inverse_event_shape(dist.event_shape))
      return self._transpose_around_bijector_fn(
          bij.inverse, arg=y,
          src_event_ndims=event_ndims, dest_event_ndims=pullback_event_ndims)
  
    def _bcast_and_reduce_logdet(self, underlying_ldj):
      # Ensure ldj is fully broadcast in the sample dims, i.e. ensure ldj has
      # full sample shape in the sample axes, before we reduce.
      batch_ndims = ps.rank_from_shape(self.distribution.batch_shape_tensor,
                                       self.distribution.batch_shape)
      extra_sample_ndims = ps.rank_from_shape(self.sample_shape)
      sample_ndims = ps.rank(underlying_ldj) - extra_sample_ndims - batch_ndims
      bcast_ldj_shape = ps.broadcast_shape(
          ps.shape(underlying_ldj),
          ps.concat([ps.ones([sample_ndims], tf.int32),
                     ps.ones([batch_ndims], tf.int32),
                     ps.reshape(self.sample_shape, shape=[-1])], axis=0))
      ldj = tf.broadcast_to(underlying_ldj, bcast_ldj_shape)
      return self._sum_fn(ldj, axis=-1 - ps.range(extra_sample_ndims))
  
    def _forward_log_det_jacobian(self, x, **kwargs):
      dist = self.distribution
      bij = self.bijector
      pullback_event_ndims = ps.rank_from_shape(
          lambda: bij.inverse_event_shape_tensor(dist.event_shape_tensor()),
          bij.inverse_event_shape(dist.event_shape))
      fn = functools.partial(bij.forward_log_det_jacobian,
                             event_ndims=pullback_event_ndims)
      underlying_ldj = self._transpose_around_bijector_fn(
          fn, arg=x, src_event_ndims=pullback_event_ndims, fn_reduces_event=True)
      return self._bcast_and_reduce_logdet(underlying_ldj)
  
    def _inverse_log_det_jacobian(self, y, **kwargs):
      event_ndims = ps.rank_from_shape(
          self.distribution.event_shape_tensor, self.distribution.event_shape)
      fn = functools.partial(self.bijector.inverse_log_det_jacobian,
                             event_ndims=event_ndims)
      underlying_ldj = self._transpose_around_bijector_fn(
          fn, arg=y, src_event_ndims=event_ndims, fn_reduces_event=True)
      return self._bcast_and_reduce_logdet(underlying_ldj)
  
  
  @kullback_leibler.RegisterKL(Sample, Sample)
  def _kl_sample(a, b, name='kl_sample'):
    """Batched KL divergence `KL(a || b)` for Sample distributions.
  
    We can leverage the fact that:
  
    ```
    KL(Sample(a) || Sample(b)) = sum(KL(a || b))
    ```
  
    where the sum is over the `sample_shape` dims.
  
    Args:
      a: Instance of `Sample` distribution.
      b: Instance of `Sample` distribution.
      name: (optional) name to use for created ops.
        Default value: `"kl_sample"`'.
  
    Returns:
      kldiv: Batchwise `KL(a || b)`.
  
    Raises:
      ValueError: If the `sample_shape` of `a` and `b` don't match.
    """
    assertions = []
    a_ss = tf.get_static_value(a.sample_shape)
    b_ss = tf.get_static_value(b.sample_shape)
    msg = '`a.sample_shape` must be identical to `b.sample_shape`.'
    if a_ss is not None and b_ss is not None:
      if not np.array_equal(a_ss, b_ss):
        raise ValueError(msg)
    elif a.validate_args or b.validate_args:
      assertions.append(assert_util.assert_equal(
          a.sample_shape, b.sample_shape, message=msg))
    with tf.control_dependencies(assertions):
      kl = kullback_leibler.kl_divergence(
          a.distribution, b.distribution, name=name)
      n = ps.reduce_prod(a.sample_shape)
      return tf.cast(x=n, dtype=kl.dtype) * kl
  
  
  @log_prob_ratio.RegisterLogProbRatio(Sample)
  def _sample_log_prob_ratio(p, x, q, y, name=None):
    """Implements `log_prob_ratio` for tfd.Sample."""
    with tf.name_scope(name or 'sample_log_prob_ratio'):
      checks = []
      if p.validate_args or q.validate_args:
        checks.append(tf.debugging.assert_equal(p.sample_shape, q.sample_shape))
      with tf.control_dependencies(checks):
        # pylint: disable=protected-access
        x, aux = p._prepare_for_underlying(x)
        y, _ = q._prepare_for_underlying(y)
        return p._finish_log_prob(
            log_prob_ratio.log_prob_ratio(p.distribution, x, q.distribution, y),
            aux)
        # pylint: enable=protected-access
  