code: # 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 Deterministic distribution class."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  
  # Dependency imports
  import six
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  
  __all__ = [
      'Deterministic',
      'VectorDeterministic',
  ]
  
  
  @six.add_metaclass(abc.ABCMeta)
  class _BaseDeterministic(distribution.Distribution):
    """Base class for Deterministic distributions."""
  
    def __init__(self,
                 loc,
                 atol=None,
                 rtol=None,
                 is_vector=False,
                 validate_args=False,
                 allow_nan_stats=True,
                 parameters=None,
                 name='_BaseDeterministic'):
      """Initialize a batch of `_BaseDeterministic` distributions.
  
      The `atol` and `rtol` parameters allow for some slack in `pmf`, `cdf`
      computations, e.g. due to floating-point error.
  
      ```
      pmf(x; loc)
        = 1, if Abs(x - loc) <= atol + rtol * Abs(loc),
        = 0, otherwise.
      ```
  
      Args:
        loc: Numeric `Tensor`.  The point (or batch of points) on which this
          distribution is supported.
        atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The absolute tolerance for comparing closeness to `loc`.
          Default is `0`.
        rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The relative tolerance for comparing closeness to `loc`.
          Default is `0`.
        is_vector:  Python `bool`.  If `True`, this is for `VectorDeterministic`,
          else `Deterministic`.
        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.
        parameters: Dict of locals to facilitate copy construction.
        name: Python `str` name prefixed to Ops created by this class.
  
      Raises:
        ValueError:  If `loc` is a scalar.
      """
      with tf.name_scope(name) as name:
        dtype = dtype_util.common_dtype([loc, atol, rtol], dtype_hint=tf.float32)
        self._loc = tensor_util.convert_nonref_to_tensor(
            loc, dtype_hint=dtype, name='loc')
        self._atol = tensor_util.convert_nonref_to_tensor(
            0 if atol is None else atol, dtype=dtype, name='atol')
        self._rtol = tensor_util.convert_nonref_to_tensor(
            0 if rtol is None else rtol, dtype=dtype, name='rtol')
        self._is_vector = is_vector
  
        super(_BaseDeterministic, self).__init__(
            dtype=self._loc.dtype,
            reparameterization_type=(
                reparameterization.FULLY_REPARAMETERIZED
                if dtype_util.is_floating(self._loc.dtype)
                else reparameterization.NOT_REPARAMETERIZED),
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            parameters=parameters,
            name=name)
  
    def _slack(self, loc):
      # Avoid using the large broadcast with self.loc if possible.
      if self.parameters['rtol'] is None:
        return self.atol
      else:
        return self.atol + self.rtol * tf.abs(loc)
  
    @property
    def loc(self):
      """Point (or batch of points) at which this distribution is supported."""
      return self._loc
  
    @property
    def atol(self):
      """Absolute tolerance for comparing points to `self.loc`."""
      return self._atol
  
    @property
    def rtol(self):
      """Relative tolerance for comparing points to `self.loc`."""
      return self._rtol
  
    def _entropy(self):
      return tf.zeros(self.batch_shape_tensor(), dtype=self.dtype)
  
    def _mean(self):
      return tf.identity(self.loc)
  
    def _variance(self):
      return tf.zeros_like(self.loc)
  
    def _mode(self):
      return self.mean()
  
    def _sample_n(self, n, seed=None):
      del seed  # unused
      loc = tf.convert_to_tensor(self.loc)
      return tf.broadcast_to(
          loc,
          tf.concat([[n], self._batch_shape_tensor(loc=loc),
                     self._event_shape_tensor(loc=loc)],
                    axis=0))
  
    def _default_event_space_bijector(self):
      return
  
    def _parameter_control_dependencies(self, is_init):
      assertions = []
  
      # In init, we can always build shape and dtype checks because
      # we assume shape doesn't change for Variable backed args.
      if is_init and self._is_vector:
        msg = 'Argument `loc` must be at least rank 1.'
        if tensorshape_util.rank(self.loc.shape) is not None:
          if tensorshape_util.rank(self.loc.shape) < 1:
            raise ValueError(msg)
        elif self.validate_args:
          assertions.append(
              assert_util.assert_rank_at_least(self.loc, 1, message=msg))
  
      if not self.validate_args:
        assert not assertions  # Should never happen
        return []
  
      if is_init != tensor_util.is_ref(self.atol):
        assertions.append(
            assert_util.assert_non_negative(
                self.atol, message='Argument "atol" must be non-negative'))
      if is_init != tensor_util.is_ref(self.rtol):
        assertions.append(
            assert_util.assert_non_negative(
                self.rtol, message='Argument "rtol" must be non-negative'))
      return assertions
  
  
  class Deterministic(_BaseDeterministic):
    """Scalar `Deterministic` distribution on the real line.
  
    The scalar `Deterministic` distribution is parameterized by a [batch] point
    `loc` on the real line.  The distribution is supported at this point only,
    and corresponds to a random variable that is constant, equal to `loc`.
  
    See [Degenerate rv](https://en.wikipedia.org/wiki/Degenerate_distribution).
  
    #### Mathematical Details
  
    The probability mass function (pmf) and cumulative distribution function (cdf)
    are
  
    ```none
    pmf(x; loc) = 1, if x == loc, else 0
    cdf(x; loc) = 1, if x >= loc, else 0
    ```
  
    #### Examples
  
    ```python
    # Initialize a single Deterministic supported at zero.
    constant = tfp.distributions.Deterministic(0.)
    constant.prob(0.)
    ==> 1.
    constant.prob(2.)
    ==> 0.
  
    # Initialize a [2, 2] batch of scalar constants.
    loc = [[0., 1.], [2., 3.]]
    x = [[0., 1.1], [1.99, 3.]]
    constant = tfp.distributions.Deterministic(loc)
    constant.prob(x)
    ==> [[1., 0.], [0., 1.]]
    ```
  
    """
  
    def __init__(self,
                 loc,
                 atol=None,
                 rtol=None,
                 validate_args=False,
                 allow_nan_stats=True,
                 name='Deterministic'):
      """Initialize a scalar `Deterministic` distribution.
  
      The `atol` and `rtol` parameters allow for some slack in `pmf`, `cdf`
      computations, e.g. due to floating-point error.
  
      ```
      pmf(x; loc)
        = 1, if Abs(x - loc) <= atol + rtol * Abs(loc),
        = 0, otherwise.
      ```
  
      Args:
        loc: Numeric `Tensor` of shape `[B1, ..., Bb]`, with `b >= 0`.
          The point (or batch of points) on which this distribution is supported.
        atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The absolute tolerance for comparing closeness to `loc`.
          Default is `0`.
        rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The relative tolerance for comparing closeness to `loc`.
          Default is `0`.
        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.
      """
      parameters = dict(locals())
      super(Deterministic, self).__init__(
          loc,
          atol=atol,
          rtol=rtol,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          name=name)
  
    @classmethod
    def _params_event_ndims(cls):
      return dict(loc=0, atol=0, rtol=0)
  
    def _batch_shape_tensor(self, loc=None):
      return tf.broadcast_dynamic_shape(
          tf.shape(self.loc if loc is None else loc),
          tf.broadcast_dynamic_shape(tf.shape(self.atol), tf.shape(self.rtol)))
  
    def _batch_shape(self):
      return tf.broadcast_static_shape(
          self.loc.shape,
          tf.broadcast_static_shape(self.atol.shape, self.rtol.shape))
  
    def _event_shape_tensor(self, loc=None):
      del loc
      return tf.constant([], dtype=tf.int32)
  
    def _event_shape(self):
      return tf.TensorShape([])
  
    def _prob(self, x):
      loc = tf.convert_to_tensor(self.loc)
      # Enforces dtype of probability to be float, when self.dtype is not.
      prob_dtype = self.dtype if dtype_util.is_floating(
          self.dtype) else tf.float32
      return tf.cast(tf.abs(x - loc) <= self._slack(loc), dtype=prob_dtype)
  
    def _cdf(self, x):
      loc = tf.identity(self.loc)
      return tf.cast(x >= loc - self._slack(loc), dtype=self.dtype)
  
  
  class VectorDeterministic(_BaseDeterministic):
    """Vector `Deterministic` distribution on `R^k`.
  
    The `VectorDeterministic` distribution is parameterized by a [batch] point
    `loc in R^k`.  The distribution is supported at this point only,
    and corresponds to a random variable that is constant, equal to `loc`.
  
    See [Degenerate rv](https://en.wikipedia.org/wiki/Degenerate_distribution).
  
    #### Mathematical Details
  
    The probability mass function (pmf) is
  
    ```none
    pmf(x; loc)
      = 1, if All[Abs(x - loc) <= atol + rtol * Abs(loc)],
      = 0, otherwise.
    ```
  
    #### Examples
  
    ```python
    tfd = tfp.distributions
  
    # Initialize a single VectorDeterministic supported at [0., 2.] in R^2.
    constant = tfd.Deterministic([0., 2.])
    constant.prob([0., 2.])
    ==> 1.
    constant.prob([0., 3.])
    ==> 0.
  
    # Initialize a [3] batch of constants on R^2.
    loc = [[0., 1.], [2., 3.], [4., 5.]]
    constant = tfd.VectorDeterministic(loc)
    constant.prob([[0., 1.], [1.9, 3.], [3.99, 5.]])
    ==> [1., 0., 0.]
    ```
  
    """
  
    def __init__(self,
                 loc,
                 atol=None,
                 rtol=None,
                 validate_args=False,
                 allow_nan_stats=True,
                 name='VectorDeterministic'):
      """Initialize a `VectorDeterministic` distribution on `R^k`, for `k >= 0`.
  
      Note that there is only one point in `R^0`, the 'point' `[]`.  So if `k = 0`
      then `self.prob([]) == 1`.
  
      The `atol` and `rtol` parameters allow for some slack in `pmf`
      computations, e.g. due to floating-point error.
  
      ```
      pmf(x; loc)
        = 1, if All[Abs(x - loc) <= atol + rtol * Abs(loc)],
        = 0, otherwise
      ```
  
      Args:
        loc: Numeric `Tensor` of shape `[B1, ..., Bb, k]`, with `b >= 0`, `k >= 0`
          The point (or batch of points) on which this distribution is supported.
        atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The absolute tolerance for comparing closeness to `loc`.
          Default is `0`.
        rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
          shape.  The relative tolerance for comparing closeness to `loc`.
          Default is `0`.
        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.
      """
      parameters = dict(locals())
      super(VectorDeterministic, self).__init__(
          loc,
          atol=atol,
          rtol=rtol,
          is_vector=True,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          name=name)
  
    @classmethod
    def _params_event_ndims(cls):
      return dict(loc=1, atol=1, rtol=1)
  
    def _batch_shape_tensor(self, loc=None):
      return tf.broadcast_dynamic_shape(
          tf.shape(self.loc if loc is None else loc),
          tf.broadcast_dynamic_shape(tf.shape(self.atol),
                                     tf.shape(self.rtol)))[:-1]
  
    def _batch_shape(self):
      return tf.broadcast_static_shape(
          self.loc.shape,
          tf.broadcast_static_shape(self.atol.shape, self.rtol.shape))[:-1]
  
    def _event_shape_tensor(self, loc=None):
      return tf.shape(self.loc if loc is None else loc)[-1:]
  
    def _event_shape(self):
      return self.loc.shape[-1:]
  
    def _prob(self, x):
      loc = tf.convert_to_tensor(self.loc)
      return tf.cast(
          tf.reduce_all(tf.abs(x - loc) <= self._slack(loc), axis=-1),
          dtype=self.dtype)
  
    def _sample_control_dependencies(self, x):
      assertions = []
      if not self.validate_args:
        return assertions
      assertions.append(assert_util.assert_rank_at_least(x, 1))
      assertions.append(assert_util.assert_equal(
          self.event_shape_tensor(), tf.gather(tf.shape(x), tf.rank(x) - 1),
          message=('Argument `x` not defined in the same space '
                   'R**k as this distribution')))
      return assertions
  
  
  @kullback_leibler.RegisterKL(_BaseDeterministic, distribution.Distribution)
  def _kl_deterministic_distribution(a, b, name=None):
    """Calculate the batched KL divergence `KL(a || b)` with `a` Deterministic.
  
    Args:
      a: instance of a Deterministic distribution object.
      b: instance of a Distribution distribution object.
      name: (optional) Name to use for created operations. Default is
        'kl_deterministic_distribution'.
  
    Returns:
      Batchwise `KL(a || b)`.
    """
    with tf.name_scope(name or 'kl_deterministic_distribution'):
      return -b.log_prob(a.loc)
   # 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.
  # ============================================================================
  """Experimental Numpy backend."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import six
  
  from tensorflow_probability.python.internal.backend.numpy import _utils as utils
  from tensorflow_probability.python.internal.backend.numpy.ops import convert_to_tensor
  from tensorflow_probability.python.internal.backend.numpy.ops import is_tensor
  from tensorflow_probability.python.internal.backend.numpy.ops import Tensor
  
  
  __all__ = [
      'Assert',
      'assert_equal',
      'assert_greater',
      'assert_greater_equal',
      'assert_integer',
      'assert_less',
      'assert_less_equal',
      'assert_near',
      'assert_negative',
      'assert_non_negative',
      'assert_non_positive',
      'assert_none_equal',
      'assert_positive',
      'assert_proper_iterable',
      'assert_rank',
      'assert_rank_at_least',
      'assert_rank_in',
      'assert_scalar',
      'check_numerics',
  ]
  
  JAX_MODE = False
  
  
  def skip_assert_for_tracers(f):
    """Function decorator that returns None if JAX tracers are detected."""
    if not JAX_MODE:
      return f
    from jax import core as jax_core  # pylint: disable=g-import-not-at-top
    def wrapped(*args, **kwargs):
      if any(isinstance(arg, jax_core.Tracer) for arg
             in args + tuple(kwargs.values())):
        print('skip assert ' + f.__name__)
        return None
      return f(*args, **kwargs)
    return wrapped
  
  
  @skip_assert_for_tracers
  def _assert_binary(
      x, y, comparator, sym, summarize=None, message=None, name=None):
    del summarize
    del name
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    if not np.all(comparator(x, y)):
      raise ValueError('Condition x {} y did not hold element-wise. {}'.format(
          sym, message or ''))
  
  
  @skip_assert_for_tracers
  def _assert_equal(x, y, summarize=None, message=None, name=None):
    del summarize
    del name
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    if not np.all(np.equal(x, y)):
      raise ValueError('Expected x == y but got {} vs {} {}'.format(
          x, y, message or ''))
  
  
  def _assert_greater(x, y, summarize=None, message=None, name=None):
    return _assert_binary(
        x, y, np.greater, '>', summarize=summarize,
        message=message, name=name)
  
  
  def _assert_less(x, y, summarize=None, message=None, name=None):
    return _assert_binary(
        x, y, np.less, '<', summarize=summarize,
        message=message, name=name)
  
  
  def _assert_greater_equal(
      x, y, summarize=None, message=None, name=None):
    return _assert_binary(
        x, y, np.greater_equal, '>=', summarize=summarize,
        message=message, name=name)
  
  
  def _assert_less_equal(
      x, y, summarize=None, message=None, name=None):
    return _assert_binary(
        x, y, np.less_equal, '<=', summarize=summarize,
        message=message, name=name)
  
  
  @skip_assert_for_tracers
  def _assert_compare_to_zero(
      x, comparator, sym, summarize=None, message=None, name=None):
    del summarize
    del name
    x = convert_to_tensor(x)
    if not np.all(comparator(x, 0)):
      raise ValueError(
          'Condition x {} 0 did not hold element-wise; got {} {}'.format(
              sym, x, message or ''))
  
  
  def _assert_positive(x, summarize=None, message=None, name=None):
    return _assert_compare_to_zero(
        x, np.greater, '>', summarize=summarize, message=message, name=name)
  
  
  def _assert_negative(x, summarize=None, message=None, name=None):
    return _assert_compare_to_zero(
        x, np.less, '<', summarize=summarize, message=message, name=name)
  
  
  def _assert_non_negative(x, summarize=None, message=None, name=None):
    return _assert_compare_to_zero(
        x, np.greater_equal, '>=',
        summarize=summarize, message=message, name=name)
  
  
  def _assert_non_positive(x, summarize=None, message=None, name=None):
    return _assert_compare_to_zero(
        x, np.less_equal, '<=', summarize=summarize, message=message, name=name)
  
  
  def _assert_rank(x, rank, message=None, name=None):  # pylint: disable=unused-argument
    return _assert_equal(x=len(np.shape(x)), y=rank, message=message)
  
  
  def _assert_scalar(*_, **__):  # pylint: disable=unused-argument
    pass
  
  
  def _assert_integer(*_, **__):  # pylint: disable=unused-argument
    pass
  
  
  @skip_assert_for_tracers
  def _assert_near(x, y, rtol=None, atol=None,
                   message=None, summarize=None, name=None):  # pylint: disable=unused-argument
    """Raises an error if abs(x - y) > atol + rtol * abs(y)."""
    del summarize
    del name
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    rtol = rtol if rtol else 10 * np.finfo(x.dtype).eps
    atol = atol if atol else 10 * np.finfo(x.dtype).eps
    if np.any(np.abs(x - y) > atol + rtol * np.abs(y)):
      raise ValueError('x = {} and y = {} are not equal to tolerance rtol = {}, '
                       'atol = {} {}'.format(x, y, rtol, atol, message or ''))
  
  
  @skip_assert_for_tracers
  def _assert_none_equal(x, y, summarize=None, message=None, name=None):
    del summarize
    del name
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    if np.any(np.equal(x, y)):
      raise ValueError('Expected x != y but got {} vs {} {}'.format(
          x, y, message or ''))
  
  
  def _assert_proper_iterable(values):
    unintentional_iterables = (Tensor, np.ndarray, bytes, six.text_type)
    if isinstance(values, unintentional_iterables):
      raise TypeError(
          'Expected argument "values" to be a "proper" iterable.  Found: %s' %
          type(values))
  
    if not hasattr(values, '__iter__'):
      raise TypeError(
          'Expected argument "values" to be iterable.  Found: %s' % type(values))
  
  
  def _assert_rank_at_least(x, rank, message=None, name=None):
    del name
    if len(x.shape) < rank:
      raise ValueError('Expected rank at least {} but got shape {} {}'.format(
          rank, x.shape, message or ''))
  
  
  def _assert_rank_in(*_, **__):  # pylint: disable=unused-argument
    pass
  
  
  # --- Begin Public Functions --------------------------------------------------
  
  
  Assert = utils.copy_docstring(  # pylint: disable=invalid-name
      'tf.debugging.Assert',
      lambda condition, data, summarize=None, name=None: None)
  
  assert_equal = utils.copy_docstring(
      'tf.debugging.assert_equal',
      _assert_equal)
  
  assert_greater = utils.copy_docstring(
      'tf.debugging.assert_greater',
      _assert_greater)
  
  assert_less = utils.copy_docstring(
      'tf.debugging.assert_less',
      _assert_less)
  
  assert_rank = utils.copy_docstring(
      'tf.debugging.assert_rank',
      _assert_rank)
  
  assert_scalar = utils.copy_docstring(
      'tf.debugging.assert_scalar',
      _assert_scalar)
  
  assert_greater_equal = utils.copy_docstring(
      'tf.debugging.assert_greater_equal',
      _assert_greater_equal)
  
  assert_integer = utils.copy_docstring(
      'tf.debugging.assert_integer',
      _assert_integer)
  
  assert_less_equal = utils.copy_docstring(
      'tf.debugging.assert_less_equal',
      _assert_less_equal)
  
  assert_near = utils.copy_docstring(
      'tf.debugging.assert_near',
      _assert_near)
  
  assert_negative = utils.copy_docstring(
      'tf.debugging.assert_negative',
      _assert_negative)
  
  assert_non_negative = utils.copy_docstring(
      'tf.debugging.assert_non_negative',
      _assert_non_negative)
  
  assert_non_positive = utils.copy_docstring(
      'tf.debugging.assert_non_positive',
      _assert_non_positive)
  
  assert_none_equal = utils.copy_docstring(
      'tf.debugging.assert_none_equal',
      _assert_none_equal)
  
  assert_positive = utils.copy_docstring(
      'tf.debugging.assert_positive',
      _assert_positive)
  
  assert_proper_iterable = utils.copy_docstring(
      'tf.debugging.assert_proper_iterable',
      _assert_proper_iterable)
  
  assert_rank_at_least = utils.copy_docstring(
      'tf.debugging.assert_rank_at_least',
      _assert_rank_at_least)
  
  assert_rank_in = utils.copy_docstring(
      'tf.debugging.assert_rank_in',
      _assert_rank_in)
  
  check_numerics = utils.copy_docstring(
      'tf.debugging.check_numerics',
      lambda x, *_, **__: x)
  
  is_numeric_tensor = utils.copy_docstring(
      'tf.debugging.is_numeric_tensor',
      lambda x: is_tensor(x) and np.issubdtype(x.dtype, np.number))
   # Copyright 2020 The TensorFlow Probability Authors. All Rights Reserved.
  # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
  # THIS FILE IS AUTO-GENERATED BY `gen_linear_operators.py`.
  # DO NOT MODIFY DIRECTLY.
  # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
  # pylint: disable=g-import-not-at-top
  # pylint: disable=g-direct-tensorflow-import
  # pylint: disable=g-bad-import-order
  # pylint: disable=unused-import
  # pylint: disable=line-too-long
  # pylint: disable=reimported
  # pylint: disable=g-bool-id-comparison
  # pylint: disable=g-statement-before-imports
  # pylint: disable=bad-continuation
  # pylint: disable=useless-import-alias
  # pylint: disable=property-with-parameters
  # pylint: disable=trailing-whitespace
  
  # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # 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.
  # ==============================================================================
  """`LinearOperator` acting like a diagonal matrix."""
  
  from __future__ import absolute_import
  from __future__ import division
  # [internal] enable type annotations
  from __future__ import print_function
  
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy import numpy_array as array_ops
  from tensorflow_probability.python.internal.backend.numpy import debugging as check_ops
  from tensorflow_probability.python.internal.backend.numpy import numpy_math as math_ops
  from tensorflow_probability.python.internal.backend.numpy import linalg_impl as linalg
  from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator
  from tensorflow_probability.python.internal.backend.numpy.gen import linear_operator_util
  # from tensorflow.python.util.tf_export import tf_export
  
  __all__ = ["LinearOperatorDiag",]
  
  
  # @tf_export("linalg.LinearOperatorDiag")
  class LinearOperatorDiag(linear_operator.LinearOperator):
    """`LinearOperator` acting like a [batch] square diagonal matrix.
  
    This operator acts like a [batch] diagonal matrix `A` with shape
    `[B1,...,Bb, N, N]` for some `b >= 0`.  The first `b` indices index a
    batch member.  For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
    an `N x N` matrix.  This matrix `A` is not materialized, but for
    purposes of broadcasting this shape will be relevant.
  
    `LinearOperatorDiag` is initialized with a (batch) vector.
  
    ```python
    # Create a 2 x 2 diagonal linear operator.
    diag = [1., -1.]
    operator = LinearOperatorDiag(diag)
  
    operator.to_dense()
    ==> [[1.,  0.]
         [0., -1.]]
  
    tensor_shape.TensorShape(operator.shape)
    ==> [2, 2]
  
    operator.log_abs_determinant()
    ==> scalar Tensor
  
    x = ... Shape [2, 4] Tensor
    operator.matmul(x)
    ==> Shape [2, 4] Tensor
  
    # Create a [2, 3] batch of 4 x 4 linear operators.
    diag = tf.random.normal(shape=[2, 3, 4])
    operator = LinearOperatorDiag(diag)
  
    # Create a shape [2, 1, 4, 2] vector.  Note that this shape is compatible
    # since the batch dimensions, [2, 1], are broadcast to
    # operator.batch_shape = [2, 3].
    y = tf.random.normal(shape=[2, 1, 4, 2])
    x = operator.solve(y)
    ==> operator.matmul(x) = y
    ```
  
    #### Shape compatibility
  
    This operator acts on [batch] matrix with compatible shape.
    `x` is a batch matrix with compatible shape for `matmul` and `solve` if
  
    ```
    tensor_shape.TensorShape(operator.shape) = [B1,...,Bb] + [N, N],  with b >= 0
    tensor_shape.TensorShape(x.shape) =   [C1,...,Cc] + [N, R],
    and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
    ```
  
    #### Performance
  
    Suppose `operator` is a `LinearOperatorDiag` of shape `[N, N]`,
    and `tensor_shape.TensorShape(x.shape) = [N, R]`.  Then
  
    * `operator.matmul(x)` involves `N * R` multiplications.
    * `operator.solve(x)` involves `N` divisions and `N * R` multiplications.
    * `operator.determinant()` involves a size `N` `reduce_prod`.
  
    If instead `operator` and `x` have shape `[B1,...,Bb, N, N]` and
    `[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
  
    #### Matrix property hints
  
    This `LinearOperator` is initialized with boolean flags of the form `is_X`,
    for `X = non_singular, self_adjoint, positive_definite, square`.
    These have the following meaning:
  
    * If `is_X == True`, callers should expect the operator to have the
      property `X`.  This is a promise that should be fulfilled, but is *not* a
      runtime assert.  For example, finite floating point precision may result
      in these promises being violated.
    * If `is_X == False`, callers should expect the operator to not have `X`.
    * If `is_X == None` (the default), callers should have no expectation either
      way.
    """
  
    def __init__(self,
                 diag,
                 is_non_singular=None,
                 is_self_adjoint=None,
                 is_positive_definite=None,
                 is_square=None,
                 name="LinearOperatorDiag"):
      r"""Initialize a `LinearOperatorDiag`.
  
      Args:
        diag:  Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`.
          The diagonal of the operator.  Allowed dtypes: `float16`, `float32`,
            `float64`, `complex64`, `complex128`.
        is_non_singular:  Expect that this operator is non-singular.
        is_self_adjoint:  Expect that this operator is equal to its hermitian
          transpose.  If `diag.dtype` is real, this is auto-set to `True`.
        is_positive_definite:  Expect that this operator is positive definite,
          meaning the quadratic form `x^H A x` has positive real part for all
          nonzero `x`.  Note that we do not require the operator to be
          self-adjoint to be positive-definite.  See:
          https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
        is_square:  Expect that this operator acts like square [batch] matrices.
        name: A name for this `LinearOperator`.
  
      Raises:
        TypeError:  If `diag.dtype` is not an allowed type.
        ValueError:  If `diag.dtype` is real, and `is_self_adjoint` is not `True`.
      """
  
      with ops.name_scope(name, values=[diag]):
        self._diag = linear_operator_util.convert_nonref_to_tensor(
            diag, name="diag")
        self._check_diag(self._diag)
  
        # Check and auto-set hints.
        if not np.issubdtype(self._diag.dtype, np.complexfloating):
          if is_self_adjoint is False:
            raise ValueError("A real diagonal operator is always self adjoint.")
          else:
            is_self_adjoint = True
  
        if is_square is False:
          raise ValueError("Only square diagonal operators currently supported.")
        is_square = True
  
        super(LinearOperatorDiag, self).__init__(
            dtype=self._diag.dtype,
            graph_parents=None,
            is_non_singular=is_non_singular,
            is_self_adjoint=is_self_adjoint,
            is_positive_definite=is_positive_definite,
            is_square=is_square,
            name=name)
        # TODO(b/143910018) Remove graph_parents in V3.
        self._set_graph_parents([self._diag])
  
    def _check_diag(self, diag):
      """Static check of diag."""
      if tensor_shape.TensorShape(diag.shape).ndims is not None and tensor_shape.TensorShape(diag.shape).ndims < 1:
        raise ValueError("Argument diag must have at least 1 dimension.  "
                         "Found: %s" % diag)
  
    def _shape(self):
      # If d_shape = [5, 3], we return [5, 3, 3].
      d_shape = tensor_shape.TensorShape(self._diag.shape)
      return d_shape.concatenate(d_shape[-1:])
  
    def _shape_tensor(self):
      d_shape = array_ops.shape(self._diag)
      k = d_shape[-1]
      return array_ops.concat((d_shape, [k]), 0)
  
    @property
    def diag(self):
      return self._diag
  
    def _assert_non_singular(self):
      return linear_operator_util.assert_no_entries_with_modulus_zero(
          self._diag,
          message="Singular operator:  Diagonal contained zero values.")
  
    def _assert_positive_definite(self):
      if np.issubdtype(self.dtype, np.complexfloating):
        message = (
            "Diagonal operator had diagonal entries with non-positive real part, "
            "thus was not positive definite.")
      else:
        message = (
            "Real diagonal operator had non-positive diagonal entries, "
            "thus was not positive definite.")
  
      return check_ops.assert_positive(
          math_ops.real(self._diag),
          message=message)
  
    def _assert_self_adjoint(self):
      return linear_operator_util.assert_zero_imag_part(
          self._diag,
          message=(
              "This diagonal operator contained non-zero imaginary values.  "
              " Thus it was not self-adjoint."))
  
    def _matmul(self, x, adjoint=False, adjoint_arg=False):
      diag_term = math_ops.conj(self._diag) if adjoint else self._diag
      x = linalg.adjoint(x) if adjoint_arg else x
      diag_mat = array_ops.expand_dims(diag_term, -1)
      return diag_mat * x
  
    def _matvec(self, x, adjoint=False):
      diag_term = math_ops.conj(self._diag) if adjoint else self._diag
      return diag_term * x
  
    def _determinant(self):
      return math_ops.reduce_prod(self._diag, axis=[-1])
  
    def _log_abs_determinant(self):
      log_det = math_ops.reduce_sum(
          math_ops.log(math_ops.abs(self._diag)), axis=[-1])
      if np.issubdtype(self.dtype, np.complexfloating):
        log_det = _ops.cast(log_det, dtype=self.dtype)
      return log_det
  
    def _solve(self, rhs, adjoint=False, adjoint_arg=False):
      diag_term = math_ops.conj(self._diag) if adjoint else self._diag
      rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
      inv_diag_mat = array_ops.expand_dims(1. / diag_term, -1)
      return rhs * inv_diag_mat
  
    def _to_dense(self):
      return _linalg.diag(self._diag)
  
    def _diag_part(self):
      return self.diag
  
    def _add_to_tensor(self, x):
      x_diag = _linalg.diag_part(x)
      new_diag = self._diag + x_diag
      return _linalg.set_diag(x, new_diag)
  
    def _eigvals(self):
      return ops.convert_to_tensor(self.diag)
  
    def _cond(self):
      abs_diag = math_ops.abs(self.diag)
      return (math_ops.reduce_max(abs_diag, axis=-1) /
              math_ops.reduce_min(abs_diag, axis=-1))
  
  import numpy as np
  from tensorflow_probability.python.internal.backend.numpy import linalg_impl as _linalg
  from tensorflow_probability.python.internal.backend.numpy import ops as _ops
  from tensorflow_probability.python.internal.backend.numpy.gen import tensor_shape
  
  from tensorflow_probability.python.internal.backend.numpy import private
  distribution_util = private.LazyLoader(
      "distribution_util", globals(),
      "tensorflow_probability.python.internal._numpy.distribution_util")
  tensorshape_util = private.LazyLoader(
     "tensorshape_util", globals(),
      "tensorflow_probability.python.internal._numpy.tensorshape_util")
  
   # 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.
  # ============================================================================
  """Trains a grammar variational auto-encoder on synthetic data.
  
  The grammar variational auto-encoder (VAE) [1] posits a generative model over
  productions from a context-free grammar, and it posits an amortized variational
  approximation for efficient posterior inference. We train the grammar VAE
  on synthetic data using the grammar from [1] (Figure 1). Note for real data
  analyses, one should implement a parser to convert examples into lists of
  production rules.
  
  This example showcases eager execution in order to train a model where data
  points have a variable number of time steps (that is, without padding). However,
  note that handling a variable number of time steps requires a batch size of 1.
  In this example, we assume data points arrive in a stream, one at a time. Such a
  setting has an unbounded maximum length which prevents padding.
  
  Summaries are written under the flag `model_dir`. Point TensorBoard to that
  directory in order to monitor progress.
  
  Example output:
  
  ```none
  Random examples from synthetic data distribution:
  222N1N21c
  1c2N2C2C12C1N
  C11C12c
  2C
  NCC
  
  Step:   0 Loss: -13.724 (0.494 sec)
  Step: 500 Loss: -0.004 (145.741 sec)
  Step: 1000 Loss: -0.000 (292.205 sec)
  Step: 1500 Loss: -0.000 (438.819 sec)
  ```
  
  #### References
  
  [1]: Matt J. Kusner, Brooks Paige, and Jose Miguel Hernandez-Lobato. Grammar
       Variational Autoencoder. In _International Conference on Machine Learning_,
       2017. https://arxiv.org/abs/1703.01925
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  import time
  
  # Dependency imports
  from absl import flags
  import six
  import tensorflow.compat.v1 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability import edward2 as ed
  
  flags.DEFINE_float("learning_rate",
                     default=1e-4,
                     help="Initial learning rate.")
  flags.DEFINE_integer("max_steps",
                       default=5000,
                       help="Number of training steps to run.")
  flags.DEFINE_integer("latent_size",
                       default=128,
                       help="Number of dimensions in the latent code.")
  flags.DEFINE_integer("num_units",
                       default=256,
                       help="Number of units in the generative model's LSTM.")
  flags.DEFINE_string("model_dir",
                      default=os.path.join(os.getenv("TEST_TMPDIR", "/tmp"),
                                           "grammar_vae/"),
                      help="Directory to put the model's fit.")
  
  FLAGS = flags.FLAGS
  
  
  class SmilesGrammar(object):
    """Context-free grammar for SMILES strings.
  
    A context-free grammar is a 4-tuple consisting of the following elements:
  
    + `nonterminal_symbols`: finite set of strings.
    + `alphabet`: finite set of strings (terminal symbols). It is disjoint from
        `nonterminal_symbols`.
    + `production_rules`: list of 2-tuples. The first and second elements of
        each tuple respectively denote the left-hand-side and right-hand-side of a
        production rule. All right-hand-sides are written as lists, since the
        number of right-hand-side symbols may be greater than 1.
    + `start_symbol`: string, a distinct nonterminal symbol.
    """
  
    @property
    def nonterminal_symbols(self):
      return {"smiles", "chain", "branched atom", "atom", "ringbond",
              "aromatic organic", "aliphatic organic", "digit"}
  
    @property
    def alphabet(self):
      return {"c", "C", "N", "1", "2"}
  
    @property
    def production_rules(self):
      return [
          ("smiles", ["chain"]),
          ("chain", ["chain", "branched atom"]),
          ("chain", ["branched atom"]),
          ("branched atom", ["atom", "ringbond"]),
          ("branched atom", ["atom"]),
          ("atom", ["aromatic organic"]),
          ("atom", ["aliphatic organic"]),
          ("ringbond", ["digit"]),
          ("aromatic organic", ["c"]),
          ("aliphatic organic", ["C"]),
          ("aliphatic organic", ["N"]),
          ("digit", ["1"]),
          ("digit", ["2"]),
      ]
  
    @property
    def start_symbol(self):
      return "smiles"
  
    def convert_to_string(self, productions):
      """Converts a sequence of productions into a string of terminal symbols.
  
      Args:
        productions: Tensor of shape [1, num_productions, num_production_rules].
          Slices along the `num_productions` dimension represent one-hot vectors.
  
      Returns:
        str that concatenates all terminal symbols from `productions`.
  
      Raises:
        ValueError: If the first production rule does not begin with
          `self.start_symbol`.
      """
      symbols = []
      for production in tf.unstack(productions, axis=1):
        lhs, rhs = self.production_rules[
            tf.argmax(input=tf.squeeze(production), axis=-1)]
        if not symbols:  # first iteration
          if lhs != self.start_symbol:
            raise ValueError("`productions` must begin with `self.start_symbol`.")
          symbols = rhs
        else:
          # Greedily unroll the nonterminal symbols based on the first occurrence
          # in a linear sequence.
          index = symbols.index(lhs)
          symbols = symbols[:index] + rhs + symbols[index + 1:]
      string = "".join(symbols)
      return string
  
    def mask(self, symbol, on_value, off_value):
      """Produces a masking tensor for (in)valid production rules.
  
      Args:
        symbol: str, a symbol in the grammar.
        on_value: Value to use for a valid production rule.
        off_value: Value to use for an invalid production rule.
  
      Returns:
        Tensor of shape [1, num_production_rules]. An element is `on_value`
        if its corresponding production rule has `symbol` on its left-hand-side;
        the element is `off_value` otherwise.
      """
      mask_values = [on_value if lhs == symbol else off_value
                     for lhs, _ in self.production_rules]
      mask_values = tf.reshape(mask_values, [1, len(self.production_rules)])
      return mask_values
  
  
  class ProbabilisticGrammar(tf.keras.Model):
    """Deep generative model over productions that follow a grammar."""
  
    def __init__(self, grammar, latent_size, num_units):
      """Constructs a probabilistic grammar.
  
      Args:
        grammar: An object representing a grammar. It has members
          `nonterminal_symbols`, `alphabet`, `production_rules`, and
          `start_symbol`, and a method `mask` determining (in)valid
          production rules given a symbol.
        latent_size: Number of dimensions in the latent code.
        num_units: Number of units in the LSTM cell.
      """
      super(ProbabilisticGrammar, self).__init__()
      self.grammar = grammar
      self.latent_size = latent_size
      self.lstm = tf.compat.v1.nn.rnn_cell.LSTMCell(num_units)
      self.output_layer = tf.keras.layers.Dense(len(grammar.production_rules))
  
    def __call__(self, *args, **kwargs):
      inputs = 0.  # fixes a dummy variable so Model can be called without inputs
      return super(ProbabilisticGrammar, self).__call__(inputs, *args, **kwargs)
  
    def call(self, inputs):
      """Runs the model forward to generate a sequence of productions.
  
      Args:
        inputs: Unused.
  
      Returns:
        productions: Tensor of shape [1, num_productions, num_production_rules].
          Slices along the `num_productions` dimension represent one-hot vectors.
      """
      del inputs  # unused
      latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),
                                              sample_shape=1,
                                              name="latent_code")
      state = self.lstm.zero_state(1, dtype=tf.float32)
      t = 0
      productions = []
      stack = [self.grammar.start_symbol]
      while stack:
        symbol = stack.pop()
        net, state = self.lstm(latent_code, state)
        logits = (self.output_layer(net) +
                  self.grammar.mask(symbol, on_value=0., off_value=-1e9))
        production = ed.OneHotCategorical(logits=logits,
                                          name="production_" + str(t))
        _, rhs = self.grammar.production_rules[tf.argmax(
            input=tf.squeeze(production), axis=-1)]
        for symbol in rhs:
          if symbol in self.grammar.nonterminal_symbols:
            stack.append(symbol)
        productions.append(production)
        t += 1
      return tf.stack(productions, axis=1)
  
  
  class ProbabilisticGrammarVariational(tf.keras.Model):
    """Amortized variational posterior for a probabilistic grammar."""
  
    def __init__(self, latent_size):
      """Constructs a variational posterior for a probabilistic grammar.
  
      Args:
        latent_size: Number of dimensions in the latent code.
      """
      super(ProbabilisticGrammarVariational, self).__init__()
      self.latent_size = latent_size
      self.encoder_net = tf.keras.Sequential([
          tf.keras.layers.Conv1D(64, 3, padding="SAME"),
          tf.keras.layers.BatchNormalization(),
          tf.keras.layers.Activation(tf.nn.elu),
          tf.keras.layers.Conv1D(128, 3, padding="SAME"),
          tf.keras.layers.BatchNormalization(),
          tf.keras.layers.Activation(tf.nn.elu),
          tf.keras.layers.Dropout(0.1),
          tf.keras.layers.GlobalAveragePooling1D(),
          tf.keras.layers.Dense(latent_size * 2, activation=None),
      ])
  
    def call(self, inputs):
      """Runs the model forward to return a stochastic encoding.
  
      Args:
        inputs: Tensor of shape [1, num_productions, num_production_rules]. It is
          a sequence of productions of length `num_productions`. Each production
          is a one-hot vector of length `num_production_rules`: it determines
          which production rule the production corresponds to.
  
      Returns:
        latent_code_posterior: A random variable capturing a sample from the
          variational distribution, of shape [1, self.latent_size].
      """
      net = self.encoder_net(tf.cast(inputs, tf.float32))
      return ed.MultivariateNormalDiag(
          loc=net[..., :self.latent_size],
          scale_diag=tf.nn.softplus(net[..., self.latent_size:]),
          name="latent_code_posterior")
  
  
  def main(argv):
    del argv  # unused
    if tf.io.gfile.exists(FLAGS.model_dir):
      tf.compat.v1.logging.warning(
          "Warning: deleting old log directory at {}".format(FLAGS.model_dir))
      tf.io.gfile.rmtree(FLAGS.model_dir)
    tf.io.gfile.makedirs(FLAGS.model_dir)
    tf.compat.v1.enable_eager_execution()
  
    grammar = SmilesGrammar()
    synthetic_data_distribution = ProbabilisticGrammar(
        grammar=grammar, latent_size=FLAGS.latent_size, num_units=FLAGS.num_units)
  
    print("Random examples from synthetic data distribution:")
    for _ in range(5):
      productions = synthetic_data_distribution()
      string = grammar.convert_to_string(productions)
      print(string)
  
    probabilistic_grammar = ProbabilisticGrammar(
        grammar=grammar, latent_size=FLAGS.latent_size, num_units=FLAGS.num_units)
    probabilistic_grammar_variational = ProbabilisticGrammarVariational(
        latent_size=FLAGS.latent_size)
  
    checkpoint = tf.train.Checkpoint(
        synthetic_data_distribution=synthetic_data_distribution,
        probabilistic_grammar=probabilistic_grammar,
        probabilistic_grammar_variational=probabilistic_grammar_variational)
    global_step = tf.compat.v1.train.get_or_create_global_step()
    optimizer = tf.compat.v1.train.AdamOptimizer(FLAGS.learning_rate)
    writer = tf.compat.v2.summary.create_file_writer(FLAGS.model_dir)
    writer.set_as_default()
  
    start_time = time.time()
    for step in range(FLAGS.max_steps):
      productions = synthetic_data_distribution()
      with tf.GradientTape() as tape:
        # Sample from amortized variational distribution and record its trace.
        with ed.tape() as variational_tape:
          _ = probabilistic_grammar_variational(productions)
  
        # Set model trace to take on the data's values and the sample from the
        # variational distribution.
        values = {"latent_code": variational_tape["latent_code_posterior"]}
        values.update({"production_" + str(t): production for t, production
                       in enumerate(tf.unstack(productions, axis=1))})
        with ed.tape() as model_tape:
          with ed.interception(ed.make_value_setter(**values)):
            _ = probabilistic_grammar()
  
        # Compute the ELBO given the variational sample, averaged over the batch
        # size and the number of time steps (number of productions). Although the
        # ELBO per data point sums over time steps, we average in order to have a
        # value that remains on the same scale across batches.
        log_likelihood = 0.
        for name, rv in six.iteritems(model_tape):
          if name.startswith("production"):
            log_likelihood += rv.distribution.log_prob(rv.value)
  
        kl = tfp.distributions.kl_divergence(
            variational_tape["latent_code_posterior"].distribution,
            model_tape["latent_code"].distribution)
  
        timesteps = tf.cast(productions.shape[1], dtype=tf.float32)
        elbo = tf.reduce_mean(input_tensor=log_likelihood - kl) / timesteps
        loss = -elbo
        with tf.compat.v2.summary.record_if(
            lambda: tf.math.equal(0, global_step % 500)):
          tf.compat.v2.summary.scalar(
              "log_likelihood",
              tf.reduce_mean(input_tensor=log_likelihood) / timesteps,
              step=global_step)
          tf.compat.v2.summary.scalar(
              "kl", tf.reduce_mean(input_tensor=kl) / timesteps, step=global_step)
          tf.compat.v2.summary.scalar("elbo", elbo, step=global_step)
  
      variables = (probabilistic_grammar.variables
                   + probabilistic_grammar_variational.variables)
      grads = tape.gradient(loss, variables)
      grads_and_vars = list(zip(grads, variables))
      optimizer.apply_gradients(grads_and_vars, global_step)
  
      if step % 500 == 0:
        duration = time.time() - start_time
        print("Step: {:>3d} Loss: {:.3f} ({:.3f} sec)".format(
            step, loss, duration))
        checkpoint.save(file_prefix=FLAGS.model_dir)
  
  if __name__ == "__main__":
    tf.compat.v1.app.run()
   # 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.
  # ============================================================================
  """Tests for the Sample distribution."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  
  from absl.testing import parameterized
  import numpy as np
  import tensorflow.compat.v2 as tf
  from tensorflow_probability.python import bijectors as tfb
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.internal import test_util
  
  
  @test_util.test_all_tf_execution_regimes
  class SampleDistributionTest(test_util.TestCase):
  
    def test_everything_scalar(self):
      s = tfd.Sample(tfd.Normal(loc=0, scale=1), 5, validate_args=True)
      x = s.sample(seed=test_util.test_seed())
      actual_lp = s.log_prob(x)
      # Sample.log_prob will reduce over event space, ie, dims [0, 2]
      # corresponding to sizes concat([[5], [2]]).
      expected_lp = tf.reduce_sum(s.distribution.log_prob(x), axis=0)
      x_, actual_lp_, expected_lp_ = self.evaluate([x, actual_lp, expected_lp])
      self.assertEqual((5,), x_.shape)
      self.assertEqual((), actual_lp_.shape)
      self.assertAllClose(expected_lp_, actual_lp_, atol=0, rtol=1e-3)
  
    def test_everything_nonscalar(self):
      s = tfd.Sample(
          tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1), [5, 4],
          validate_args=True)
      x = s.sample([6, 1], seed=test_util.test_seed())
      actual_lp = s.log_prob(x)
      # Sample.log_prob will reduce over event space, ie, dims [2, 3, 5]
      # corresponding to sizes concat([[5, 4], [2]]).
      expected_lp = tf.reduce_sum(
          s.distribution.log_prob(tf.transpose(a=x, perm=[0, 1, 3, 4, 2, 5])),
          axis=[2, 3])
      x_, actual_lp_, expected_lp_ = self.evaluate([x, actual_lp, expected_lp])
      self.assertEqual((6, 1, 3, 5, 4, 2), x_.shape)
      self.assertEqual((6, 1, 3), actual_lp_.shape)
      self.assertAllClose(expected_lp_, actual_lp_, atol=0, rtol=1e-3)
  
    def test_mixed_scalar(self):
      s = tfd.Sample(tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1),
                     3, validate_args=False)
      x = s.sample(4, seed=test_util.test_seed())
      lp = s.log_prob(x)
      self.assertEqual((4, 3, 2), x.shape)
      self.assertEqual((4,), lp.shape)
  
    def test_kl_divergence(self):
      q_scale = 2.
      p = tfd.Sample(
          tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1), [5, 4],
          validate_args=True)
      q = tfd.Sample(
          tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=2.), 1), [5, 4],
          validate_args=True)
      actual_kl = tfd.kl_divergence(p, q)
      expected_kl = ((5 * 4) *
                     (0.5 * q_scale**-2. - 0.5 + np.log(q_scale)) *  # Actual KL.
                     np.ones([3]) * 2)  # Batch, events.
      self.assertAllClose(expected_kl, self.evaluate(actual_kl))
  
    def test_transformed_affine(self):
      sample_shape = 3
      mvn = tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1)
      aff = tfb.Affine(scale_tril=[[0.75, 0.],
                                   [0.05, 0.5]])
  
      def expected_lp(y):
        x = aff.inverse(y)  # Ie, tf.random.normal([4, 3, 2])
        fldj = aff.forward_log_det_jacobian(x, event_ndims=1)
        return tf.reduce_sum(mvn.log_prob(x) - fldj, axis=1)
  
      # Transform a Sample.
      d = tfd.TransformedDistribution(
          tfd.Sample(mvn, sample_shape, validate_args=True),
          bijector=aff)
      y = d.sample(4, seed=test_util.test_seed())
      actual_lp = d.log_prob(y)
      self.assertAllEqual((4,) + (sample_shape,) + (2,), y.shape)
      self.assertAllEqual((4,), actual_lp.shape)
      self.assertAllClose(
          *self.evaluate([expected_lp(y), actual_lp]),
          atol=0., rtol=1e-3)
  
      # Sample a Transform.
      d = tfd.Sample(
          tfd.TransformedDistribution(mvn, bijector=aff),
          sample_shape,
          validate_args=True)
      y = d.sample(4, seed=test_util.test_seed())
      actual_lp = d.log_prob(y)
      self.assertAllEqual((4,) + (sample_shape,) + (2,), y.shape)
      self.assertAllEqual((4,), actual_lp.shape)
      self.assertAllClose(
          *self.evaluate([expected_lp(y), actual_lp]),
          atol=0., rtol=1e-3)
  
    def test_transformed_exp(self):
      sample_shape = 3
      mvn = tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1)
      exp = tfb.Exp()
  
      def expected_lp(y):
        x = exp.inverse(y)  # Ie, tf.random.normal([4, 3, 2])
        fldj = exp.forward_log_det_jacobian(x, event_ndims=1)
        return tf.reduce_sum(mvn.log_prob(x) - fldj, axis=1)
  
      # Transform a Sample.
      d = tfd.TransformedDistribution(
          tfd.Sample(mvn, sample_shape, validate_args=True),
          bijector=exp)
      y = d.sample(4, seed=test_util.test_seed())
      actual_lp = d.log_prob(y)
      self.assertAllEqual((4,) + (sample_shape,) + (2,), y.shape)
      self.assertAllEqual((4,), actual_lp.shape)
      # If `TransformedDistribution` didn't scale the jacobian by
      # `_sample_distribution_size`, then `scale_fldj` would need to be `False`.
      self.assertAllClose(
          *self.evaluate([expected_lp(y), actual_lp]),
          atol=0., rtol=1e-3)
  
      # Sample a Transform.
      d = tfd.Sample(
          tfd.TransformedDistribution(mvn, bijector=exp),
          sample_shape,
          validate_args=True)
      y = d.sample(4, seed=test_util.test_seed())
      actual_lp = d.log_prob(y)
      self.assertAllEqual((4,) + (sample_shape,) + (2,), y.shape)
      self.assertAllEqual((4,), actual_lp.shape)
      # Regardless of whether `TransformedDistribution` scales the jacobian by
      # `_sample_distribution_size`, `scale_fldj` is `True`.
      self.assertAllClose(
          *self.evaluate([expected_lp(y), actual_lp]),
          atol=0., rtol=1e-3)
  
    @parameterized.parameters(
        'mean',
        'stddev',
        'variance',
        'mode',
    )
    def test_summary_statistic(self, attr):
      sample_shape = [5, 4]
      mvn = tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1)
      d = tfd.Sample(mvn, sample_shape, validate_args=True)
      self.assertEqual((3,), d.batch_shape)
      expected_stat = (
          getattr(mvn, attr)()[:, tf.newaxis, tf.newaxis, :] *
          tf.ones([3, 5, 4, 2]))
      actual_stat = getattr(d, attr)()
      self.assertAllEqual(*self.evaluate([expected_stat, actual_stat]))
  
    def test_entropy(self):
      sample_shape = [3, 4]
      mvn = tfd.Independent(tfd.Normal(loc=0, scale=[[0.25, 0.5]]), 1)
      d = tfd.Sample(mvn, sample_shape, validate_args=True)
      expected_entropy = 12 * tf.reduce_sum(mvn.distribution.entropy(), axis=-1)
      actual_entropy = d.entropy()
      self.assertAllEqual(*self.evaluate([expected_entropy, actual_entropy]))
  
    @test_util.tf_tape_safety_test
    def test_gradients_through_params(self):
      loc = tf.Variable(tf.zeros([4, 5, 3]), shape=tf.TensorShape(None))
      scale = tf.Variable(tf.ones([]), shape=tf.TensorShape(None))
      # In real life, you'd really always want `sample_shape` to be
      # `trainable=False`.
      sample_shape = tf.Variable([1, 2], shape=tf.TensorShape(None))
      dist = tfd.Sample(
          tfd.Independent(tfd.Logistic(loc=loc, scale=scale),
                          reinterpreted_batch_ndims=1),
          sample_shape=sample_shape,
          validate_args=True)
      with tf.GradientTape() as tape:
        loss = -dist.log_prob(0.)
      self.assertLen(dist.trainable_variables, 3)
      grad = tape.gradient(loss, [loc, scale, sample_shape])
      self.assertAllNotNone(grad[:-1])
      self.assertIs(grad[-1], None)
  
    @test_util.tf_tape_safety_test
    def test_variable_shape_change(self):
      loc = tf.Variable(tf.zeros([4, 5, 3]), shape=tf.TensorShape(None))
      scale = tf.Variable(tf.ones([]), shape=tf.TensorShape(None))
      # In real life, you'd really always want `sample_shape` to be
      # `trainable=False`.
      sample_shape = tf.Variable([1, 2], shape=tf.TensorShape(None))
      dist = tfd.Sample(
          tfd.Independent(tfd.Logistic(loc=loc, scale=scale),
                          reinterpreted_batch_ndims=1),
          sample_shape=sample_shape,
          validate_args=True)
      self.evaluate([v.initializer for v in dist.trainable_variables])
  
      x = dist.mean()
      y = dist.sample([7, 2], seed=test_util.test_seed())
      loss_x = -dist.log_prob(x)
      loss_0 = -dist.log_prob(0.)
      batch_shape = dist.batch_shape_tensor()
      event_shape = dist.event_shape_tensor()
      [x_, y_, loss_x_, loss_0_, batch_shape_, event_shape_] = self.evaluate([
          x, y, loss_x, loss_0, batch_shape, event_shape])
      self.assertAllEqual([4, 5, 1, 2, 3], x_.shape)
      self.assertAllEqual([7, 2, 4, 5, 1, 2, 3], y_.shape)
      self.assertAllEqual([4, 5], loss_x_.shape)
      self.assertAllEqual([4, 5], loss_0_.shape)
      self.assertAllEqual([4, 5], batch_shape_)
      self.assertAllEqual([1, 2, 3], event_shape_)
      self.assertLen(dist.trainable_variables, 3)
  
      with tf.control_dependencies([
          loc.assign(tf.zeros([])),
          scale.assign(tf.ones([3, 1, 2])),
          sample_shape.assign(6),
      ]):
        x = dist.mean()
        y = dist.sample([7, 2], seed=test_util.test_seed())
        loss_x = -dist.log_prob(x)
        loss_0 = -dist.log_prob(0.)
        batch_shape = dist.batch_shape_tensor()
        event_shape = dist.event_shape_tensor()
      [x_, y_, loss_x_, loss_0_, batch_shape_, event_shape_] = self.evaluate([
          x, y, loss_x, loss_0, batch_shape, event_shape])
      self.assertAllEqual([3, 1, 6, 2], x_.shape)
      self.assertAllEqual([7, 2, 3, 1, 6, 2], y_.shape)
      self.assertAllEqual([3, 1], loss_x_.shape)
      self.assertAllEqual([3, 1], loss_0_.shape)
      self.assertAllEqual([3, 1], batch_shape_)
      self.assertAllEqual([6, 2], event_shape_)
      self.assertLen(dist.trainable_variables, 3)
  
    def test_variable_sample_shape_exception(self):
      loc = tf.Variable(tf.zeros([4, 5, 3]), shape=tf.TensorShape(None))
      scale = tf.Variable(tf.ones([]), shape=tf.TensorShape(None))
      sample_shape = tf.Variable([[1, 2]], shape=tf.TensorShape(None))
      with self.assertRaisesWithPredicateMatch(
          Exception,
          'Argument `sample_shape` must be either a scalar or a vector.'):
        dist = tfd.Sample(
            tfd.Independent(tfd.Logistic(loc=loc, scale=scale),
                            reinterpreted_batch_ndims=1),
            sample_shape=sample_shape,
            validate_args=True)
        self.evaluate([v.initializer for v in dist.trainable_variables])
        self.evaluate(dist.mean())
  
  
  if __name__ == '__main__':
    tf.test.main()
  
apis: tensorflow.compat.v2.shape tensorflow.compat.v2.convert_to_tensor tensorflow.compat.v2.identity tensorflow.compat.v2.TensorShape tensorflow.compat.v2.abs tensorflow.compat.v2.name_scope tensorflow.compat.v2.broadcast_static_shape tensorflow.compat.v2.constant tensorflow.compat.v2.rank tensorflow.compat.v2.zeros_like
apis: numpy.equal numpy.issubdtype numpy.abs numpy.shape numpy.finfo
apis: numpy.issubdtype
apis: tensorflow.compat.v1.compat.v1.nn.rnn_cell.LSTMCell tensorflow.compat.v1.nn.softplus tensorflow.compat.v1.unstack tensorflow.compat.v1.keras.layers.Activation tensorflow.compat.v1.keras.layers.Dense tensorflow.compat.v1.io.gfile.rmtree tensorflow.compat.v1.compat.v1.train.get_or_create_global_step tensorflow.compat.v1.keras.layers.Conv1D tensorflow.compat.v1.io.gfile.exists tensorflow.compat.v1.reduce_mean tensorflow.compat.v1.train.Checkpoint tensorflow.compat.v1.io.gfile.makedirs tensorflow.compat.v1.cast tensorflow.compat.v1.stack tensorflow.compat.v1.compat.v1.train.AdamOptimizer tensorflow.compat.v1.compat.v2.summary.create_file_writer tensorflow.compat.v1.keras.layers.BatchNormalization tensorflow.compat.v1.math.equal tensorflow.compat.v1.compat.v2.summary.scalar tensorflow.compat.v1.compat.v1.enable_eager_execution tensorflow.compat.v1.keras.layers.GlobalAveragePooling1D tensorflow.compat.v1.compat.v1.app.run tensorflow.compat.v1.zeros tensorflow.compat.v1.GradientTape tensorflow.compat.v1.keras.layers.Dropout tensorflow.compat.v1.squeeze
apis: tensorflow.compat.v2.transpose numpy.ones tensorflow.compat.v2.zeros tensorflow.compat.v2.test.main tensorflow.compat.v2.GradientTape numpy.log tensorflow.compat.v2.ones tensorflow.compat.v2.TensorShape