# 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 Chi distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import invert as invert_bijector
from tensorflow_probability.python.bijectors import square as square_bijector
from tensorflow_probability.python.distributions import chi2
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import transformed_distribution
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util

class Chi(transformed_distribution.TransformedDistribution):
  """Chi distribution.

  The Chi distribution is defined over nonnegative real numbers and uses a
  degrees of freedom ("df") parameter.

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(x; df, x >= 0) = x**(df - 1) exp(-0.5 x**2) / Z
  Z = 2**(0.5 df - 1) Gamma(0.5 df)
  ```

  where:

  * `df` denotes the degrees of freedom,
  * `Z` is the normalization constant, and,
  * `Gamma` is the [gamma function](
    https://en.wikipedia.org/wiki/Gamma_function).

  The Chi distribution is a transformation of the Chi2 distribution; it is the
  distribution of the positive square root of a variable obeying a Chi2
  distribution.

  """

  def __init__(self,
               df,
               validate_args=False,
               allow_nan_stats=True,
               name="Chi"):
    """Construct Chi 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.
        Default value: `'Chi'`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      df = tf.convert_to_tensor(
          value=df,
          name="df",
          dtype=dtype_util.common_dtype([df], preferred_dtype=tf.float32))
      validation_assertions = (
          [assert_util.assert_positive(df)] if validate_args else [])
      with tf.control_dependencies(validation_assertions):
        self._df = tf.identity(df, name="df")

      super(Chi, self).__init__(
          distribution=chi2.Chi2(df=self._df,
                                 validate_args=validate_args,
                                 allow_nan_stats=allow_nan_stats),
          bijector=invert_bijector.Invert(square_bijector.Square()),
          parameters=parameters,
          name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(df=0)

  @property
  def df(self):
    """Distribution parameter for degrees of freedom."""
    return self._df

  def _mean(self):
    return np.sqrt(2) * tf.exp(
        tf.math.lgamma(0.5 * (self.df + 1)) - tf.math.lgamma(0.5 * self.df))

  def _variance(self):
    return self.df - tf.square(self._mean())

  def _entropy(self):
    return (tf.math.lgamma(self.df / 2) + 0.5 *
            (self.df - np.log(2) -
             (self.df - 1) * tf.math.digamma(0.5 * self.df)))

@kullback_leibler.RegisterKL(Chi, Chi)
def _kl_chi_chi(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Chi.

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_chi_chi"):
    # Consistent with
    # https://mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 118
    # The paper introduces an additional scaling parameter; setting that
    # parameter to 1 and simplifying yields the expression we use here.
    return (0.5 * tf.math.digamma(0.5 * a.df) * (a.df - b.df) +
            tf.math.lgamma(0.5 * b.df) - tf.math.lgamma(0.5 * a.df))


# 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.
# ============================================================================
"""Hessian-type proximal descent optimizer.

This optimizer uses proximal gradient descent and a step size dependent on the
Hessian to efficiently minimize a convex loss function with L1 and L2
regularization.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.math.linalg import sparse_or_dense_matvecmul
from tensorflow_probability.python.math.numeric import soft_threshold

__all__ = [
    'minimize',
    'minimize_one_step',
]

def _reduce_ignoring_nones(fn, args):
  result = None
  for x in args:
    if x is not None:
      result = (x if result is None else fn(result, x))
  return result

def _reduce_exiting_on_none(fn, args):
  args = list(args)
  if any(x is None for x in args):
    return None
  return _reduce_ignoring_nones(fn, args)

def _add_ignoring_nones(*args):
  return _reduce_ignoring_nones(lambda a, b: a + b, args)

def _mul_ignoring_nones(*args):
  return _reduce_ignoring_nones(lambda a, b: a * b, args)

def _mul_or_none(*args):
  return _reduce_exiting_on_none(lambda a, b: a * b, args)

def _get_shape(x, out_type=tf.int32):
  # Return the shape of a Tensor or a SparseTensor as an np.array if its shape
  # is known statically. Otherwise return a Tensor representing the shape.
  if x.shape.is_fully_defined():
    return np.array(x.shape.as_list(), dtype=out_type.as_numpy_dtype)
  return tf.shape(input=x, out_type=out_type)

def _sparse_or_dense_matmul_onehot(sparse_or_dense_matrix, col_index):
  """Returns a (dense) column of a Tensor or SparseTensor.

  Args:
    sparse_or_dense_matrix: matrix-shaped, `float` `Tensor` or `SparseTensor`.
    col_index: scalar, `int` `Tensor` representing the index of the desired
      column.

  Returns:
    column: vector-shaped, `float` `Tensor` with the same dtype as
      `sparse_or_dense_matrix`, representing the `col_index`th column of
      `sparse_or_dense_matrix`.
  """
  if isinstance(sparse_or_dense_matrix,
                (tf.SparseTensor, tf.compat.v1.SparseTensorValue)):
    # TODO(b/111924846): Implement better (ideally in a way that allows us to
    # eliminate the `num_rows` arg, if possible).
    num_rows = _get_shape(sparse_or_dense_matrix)[-2]
    batch_shape = _get_shape(sparse_or_dense_matrix)[:-2]
    slice_start = tf.concat([tf.zeros_like(batch_shape), [0, col_index]],
                            axis=0)
    slice_size = tf.concat([batch_shape, [num_rows, 1]], axis=0)
    # We momentarily lose static shape information in tf.sparse_slice. However
    # we regain it in the following tf.reshape.
    sparse_slice = tf.sparse.slice(sparse_or_dense_matrix,
                                   tf.cast(slice_start, tf.int64),
                                   tf.cast(slice_size, tf.int64))

    output_shape = tf.concat([batch_shape, [num_rows]], axis=0)
    return tf.reshape(tf.sparse.to_dense(sparse_slice), output_shape)
  else:
    return tf.gather(sparse_or_dense_matrix, col_index, axis=-1)

def _one_hot_like(x, indices, on_value=None):
  output_dtype = x.dtype.base_dtype
  if tf.compat.dimension_value(x.shape[-1]) is None:
    depth = tf.shape(input=x)[-1]
  else:
    depth = tf.compat.dimension_value(x.shape[-1])
  if on_value is not None:
    on_value = tf.cast(on_value, output_dtype)
  return tf.one_hot(indices, depth=depth, on_value=on_value, dtype=output_dtype)

def minimize_one_step(gradient_unregularized_loss,
                      hessian_unregularized_loss_outer,
                      hessian_unregularized_loss_middle,
                      x_start,
                      tolerance,
                      l1_regularizer,
                      l2_regularizer=None,
                      maximum_full_sweeps=1,
                      learning_rate=None,
                      name=None):
  """One step of (the outer loop of) the minimization algorithm.

  This function returns a new value of `x`, equal to `x_start + x_update`.  The
  increment `x_update in R^n` is computed by a coordinate descent method, that
  is, by a loop in which each iteration updates exactly one coordinate of
  `x_update`.  (Some updates may leave the value of the coordinate unchanged.)

  The particular update method used is to apply an L1-based proximity operator,
  "soft threshold", whose fixed point `x_update_fix` is the desired minimum

  ```none
  x_update_fix = argmin{
      Loss(x_start + x_update')
        + l1_regularizer * ||x_start + x_update'||_1
        + l2_regularizer * ||x_start + x_update'||_2**2
      : x_update' }
  ```

  where in each iteration `x_update'` is constrained to have at most one nonzero
  coordinate.

  This update method preserves sparsity, i.e., tends to find sparse solutions if
  `x_start` is sparse.  Additionally, the choice of step size is based on
  curvature (Hessian), which significantly speeds up convergence.

  This algorithm assumes that `Loss` is convex, at least in a region surrounding
  the optimum.  (If `l2_regularizer > 0`, then only weak convexity is needed.)

  Args:
    gradient_unregularized_loss: (Batch of) `Tensor` with the same shape and
      dtype as `x_start` representing the gradient, evaluated at `x_start`, of
      the unregularized loss function (denoted `Loss` above).  (In all current
      use cases, `Loss` is the negative log likelihood.)
    hessian_unregularized_loss_outer: (Batch of) `Tensor` or `SparseTensor`
      having the same dtype as `x_start`, and shape `[N, n]` where `x_start` has
      shape `[n]`, satisfying the property
      `Transpose(hessian_unregularized_loss_outer)
      @ diag(hessian_unregularized_loss_middle)
      @ hessian_unregularized_loss_inner
      = (approximation of) Hessian matrix of Loss, evaluated at x_start`.
    hessian_unregularized_loss_middle: (Batch of) vector-shaped `Tensor` having
      the same dtype as `x_start`, and shape `[N]` where
      `hessian_unregularized_loss_outer` has shape `[N, n]`, satisfying the
      property
      `Transpose(hessian_unregularized_loss_outer)
      @ diag(hessian_unregularized_loss_middle)
      @ hessian_unregularized_loss_inner
      = (approximation of) Hessian matrix of Loss, evaluated at x_start`.
    x_start: (Batch of) vector-shaped, `float` `Tensor` representing the current
      value of the argument to the Loss function.
    tolerance: scalar, `float` `Tensor` representing the convergence threshold.
      The optimization step will terminate early, returning its current value of
      `x_start + x_update`, once the following condition is met:
      `||x_update_end - x_update_start||_2 / (1 + ||x_start||_2)
      < sqrt(tolerance)`,
      where `x_update_end` is the value of `x_update` at the end of a sweep and
      `x_update_start` is the value of `x_update` at the beginning of that
      sweep.
    l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
      regularization term (see equation above).  If L1 regularization is not
      required, then `tfp.glm.fit_one_step` is preferable.
    l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
      regularization term (see equation above).
      Default value: `None` (i.e., no L2 regularization).
    maximum_full_sweeps: Python integer specifying maximum number of sweeps to
      run.  A "sweep" consists of an iteration of coordinate descent on each
      coordinate. After this many sweeps, the algorithm will terminate even if
      convergence has not been reached.
      Default value: `1`.
    learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
      used to dampen the proximal gradient descent steps.
      Default value: `None` (i.e., factor is conceptually `1`).
    name: Python string representing the name of the TensorFlow operation.
      The default name is `"minimize_one_step"`.

  Returns:
    x: (Batch of) `Tensor` having the same shape and dtype as `x_start`,
      representing the updated value of `x`, that is, `x_start + x_update`.
    is_converged: scalar, `bool` `Tensor` indicating whether convergence
      occurred across all batches within the specified number of sweeps.
    iter: scalar, `int` `Tensor` representing the actual number of coordinate
      updates made (before achieving convergence).  Since each sweep consists of
      `tf.size(x_start)` iterations, the maximum number of updates is
      `maximum_full_sweeps * tf.size(x_start)`.

  #### References

  [1]: Jerome Friedman, Trevor Hastie and Rob Tibshirani. Regularization Paths
       for Generalized Linear Models via Coordinate Descent. _Journal of
       Statistical Software_, 33(1), 2010.
       https://www.jstatsoft.org/article/view/v033i01/v33i01.pdf

  [2]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for
       L1-regularized Logistic Regression. _Journal of Machine Learning
       Research_, 13, 2012.
       http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf
  """
  graph_deps = [
      gradient_unregularized_loss,
      hessian_unregularized_loss_outer,
      hessian_unregularized_loss_middle,
      x_start,
      l1_regularizer,
      l2_regularizer,
      maximum_full_sweeps,
      tolerance,
      learning_rate,
  ]
  with tf.compat.v1.name_scope(name, 'minimize_one_step', graph_deps):
    x_shape = _get_shape(x_start)
    batch_shape = x_shape[:-1]
    dims = x_shape[-1]

    def _hessian_diag_elt_with_l2(coord):  # pylint: disable=missing-docstring
      # Returns the (coord, coord) entry of
      #
      #   Hessian(UnregularizedLoss(x) + l2_regularizer * ||x||_2**2)
      #
      # evaluated at x = x_start.
      inner_square = tf.reduce_sum(
          input_tensor=_sparse_or_dense_matmul_onehot(
              hessian_unregularized_loss_outer, coord)**2,
          axis=-1)
      unregularized_component = (
          hessian_unregularized_loss_middle[..., coord] * inner_square)
      l2_component = _mul_or_none(2., l2_regularizer)
      return _add_ignoring_nones(unregularized_component, l2_component)

    grad_loss_with_l2 = _add_ignoring_nones(
        gradient_unregularized_loss, _mul_or_none(2., l2_regularizer, x_start))

    # We define `x_update_diff_norm_sq_convergence_threshold` such that the
    # convergence condition
    #     ||x_update_end - x_update_start||_2 / (1 + ||x_start||_2)
    #     < sqrt(tolerance)
    # is equivalent to
    #     ||x_update_end - x_update_start||_2**2
    #     < x_update_diff_norm_sq_convergence_threshold.
    x_update_diff_norm_sq_convergence_threshold = (
        tolerance * (1. + tf.norm(tensor=x_start, ord=2, axis=-1))**2)

    # Reshape update vectors so that the coordinate sweeps happen along the
    # first dimension. This is so that we can use tensor_scatter_update to make
    # sparse updates along the first axis without copying the Tensor.
    # TODO(b/118789120): Switch to something like tf.tensor_scatter_nd_add if
    # or when it exists.
    update_shape = tf.concat([[dims], batch_shape], axis=-1)

    def _loop_cond(iter_, x_update_diff_norm_sq, x_update,
                   hess_matmul_x_update):
      del x_update
      del hess_matmul_x_update
      sweep_complete = (iter_ > 0) & tf.equal(iter_ % dims, 0)
      small_delta = (
          x_update_diff_norm_sq < x_update_diff_norm_sq_convergence_threshold)
      converged = sweep_complete & small_delta
      allowed_more_iterations = iter_ < maximum_full_sweeps * dims
      return allowed_more_iterations & tf.reduce_any(input_tensor=~converged)

    def _loop_body(  # pylint: disable=missing-docstring
        iter_, x_update_diff_norm_sq, x_update, hess_matmul_x_update):
      # Inner loop of the minimizer.
      #
      # This loop updates a single coordinate of x_update.  Ideally, an
      # iteration of this loop would set
      #
      #   x_update[j] += argmin{ LocalLoss(x_update + z*e_j) : z in R }
      #
      # where
      #
      #   LocalLoss(x_update')
      #     = LocalLossSmoothComponent(x_update')
      #         + l1_regularizer * (||x_start + x_update'||_1 -
      #                             ||x_start + x_update||_1)
      #    := (UnregularizedLoss(x_start + x_update') -
      #        UnregularizedLoss(x_start + x_update)
      #         + l2_regularizer * (||x_start + x_update'||_2**2 -
      #                             ||x_start + x_update||_2**2)
      #         + l1_regularizer * (||x_start + x_update'||_1 -
      #                             ||x_start + x_update||_1)
      #
      # In this algorithm approximate the above argmin using (univariate)
      # proximal gradient descent:
      #
      # (*)  x_update[j] = prox_{t * l1_regularizer * L1}(
      #                 x_update[j] -
      #                 t * d/dz|z=0 UnivariateLocalLossSmoothComponent(z))
      #
      # where
      #
      #   UnivariateLocalLossSmoothComponent(z)
      #       := LocalLossSmoothComponent(x_update + z*e_j)
      #
      # and we approximate
      #
      #       d/dz UnivariateLocalLossSmoothComponent(z)
      #     = grad LocalLossSmoothComponent(x_update))[j]
      #    ~= (grad LossSmoothComponent(x_start)
      #         + x_update matmul HessianOfLossSmoothComponent(x_start))[j].
      #
      # To choose the parameter t, we squint and pretend that the inner term of
      # (*) is a Newton update as if we were using Newton's method to minimize
      # UnivariateLocalLossSmoothComponent.  That is, we choose t such that
      #
      #   -t * d/dz ULLSC = -learning_rate * (d/dz ULLSC) / (d^2/dz^2 ULLSC)
      #
      # at z=0.  Hence
      #
      #   t = learning_rate / (d^2/dz^2|z=0 ULLSC)
      #     = learning_rate / HessianOfLossSmoothComponent(
      #                           x_start + x_update)[j,j]
      #    ~= learning_rate / HessianOfLossSmoothComponent(
      #                           x_start)[j,j]
      #
      # The above approximation is equivalent to assuming that
      # HessianOfUnregularizedLoss is constant, i.e., ignoring third-order
      # effects.
      #
      # Note that because LossSmoothComponent is (assumed to be) convex, t is
      # positive.

      # In above notation, coord = j.
      coord = iter_ % dims
      # x_update_diff_norm_sq := ||x_update_end - x_update_start||_2**2,
      # computed incrementally, where x_update_end and x_update_start are as
      # defined in the convergence criteria.  Accordingly, we reset
      # x_update_diff_norm_sq to zero at the beginning of each sweep.
      x_update_diff_norm_sq = tf.where(
          tf.equal(coord, 0), tf.zeros_like(x_update_diff_norm_sq),
          x_update_diff_norm_sq)

      # Recall that x_update and hess_matmul_x_update has the rightmost
      # dimension transposed to the leftmost dimension.
      w_old = x_start[..., coord] + x_update[coord, ...]
      # This is the coordinatewise Newton update if no L1 regularization.
      # In above notation, newton_step = -t * (approximation of d/dz|z=0 ULLSC).
      second_deriv = _hessian_diag_elt_with_l2(coord)
      newton_step = -_mul_ignoring_nones(  # pylint: disable=invalid-unary-operand-type
          learning_rate, grad_loss_with_l2[..., coord] +
          hess_matmul_x_update[coord, ...]) / second_deriv

      # Applying the soft-threshold operator accounts for L1 regularization.
      # In above notation, delta =
      #     prox_{t*l1_regularizer*L1}(w_old + newton_step) - w_old.
      delta = (
          soft_threshold(
              w_old + newton_step,
              _mul_ignoring_nones(learning_rate, l1_regularizer) / second_deriv)
          - w_old)

      def _do_update(x_update_diff_norm_sq, x_update, hess_matmul_x_update):  # pylint: disable=missing-docstring
        hessian_column_with_l2 = sparse_or_dense_matvecmul(
            hessian_unregularized_loss_outer,
            hessian_unregularized_loss_middle * _sparse_or_dense_matmul_onehot(
                hessian_unregularized_loss_outer, coord),
            adjoint_a=True)

        if l2_regularizer is not None:
          hessian_column_with_l2 += _one_hot_like(
              hessian_column_with_l2, coord, on_value=2. * l2_regularizer)

        # Move the batch dimensions of `hessian_column_with_l2` to rightmost in
        # order to conform to `hess_matmul_x_update`.
        n = tf.rank(hessian_column_with_l2)
        perm = tf.roll(tf.range(n), shift=1, axis=0)
        hessian_column_with_l2 = tf.transpose(
            a=hessian_column_with_l2, perm=perm)

        # Update the entire batch at `coord` even if `delta` may be 0 at some
        # batch coordinates. In those cases, adding `delta` is a no-op.
        x_update = tf.tensor_scatter_nd_add(x_update, [[coord]], [delta])

        with tf.control_dependencies([x_update]):
          x_update_diff_norm_sq_ = x_update_diff_norm_sq + delta**2
          hess_matmul_x_update_ = (
              hess_matmul_x_update + delta * hessian_column_with_l2)

          # Hint that loop vars retain the same shape.
          x_update_diff_norm_sq_.set_shape(
              x_update_diff_norm_sq_.shape.merge_with(
                  x_update_diff_norm_sq.shape))
          hess_matmul_x_update_.set_shape(
              hess_matmul_x_update_.shape.merge_with(
                  hess_matmul_x_update.shape))

          return [x_update_diff_norm_sq_, x_update, hess_matmul_x_update_]

      inputs_to_update = [x_update_diff_norm_sq, x_update, hess_matmul_x_update]
      return [iter_ + 1] + prefer_static.cond(
          # Note on why checking delta (a difference of floats) for equality to
          # zero is ok:
          #
          # First of all, x - x == 0 in floating point -- see
          # https://stackoverflow.com/a/2686671
          #
          # Delta will conceptually equal zero when one of the following holds:
          # (i)   |w_old + newton_step| <= threshold and w_old == 0
          # (ii)  |w_old + newton_step| > threshold and
          #       w_old + newton_step - sign(w_old + newton_step) * threshold
          #          == w_old
          #
          # In case (i) comparing delta to zero is fine.
          #
          # In case (ii), newton_step conceptually equals
          #     sign(w_old + newton_step) * threshold.
          # Also remember
          #     threshold = -newton_step / (approximation of d/dz|z=0 ULLSC).
          # So (i) happens when
          #     (approximation of d/dz|z=0 ULLSC) == -sign(w_old + newton_step).
          # If we did not require LossSmoothComponent to be strictly convex,
          # then this could actually happen a non-negligible amount of the time,
          # e.g. if the loss function is piecewise linear and one of the pieces
          # has slope 1.  But since LossSmoothComponent is strictly convex, (i)
          # should not systematically happen.
          tf.reduce_all(input_tensor=tf.equal(delta, 0.)),
          lambda: inputs_to_update,
          lambda: _do_update(*inputs_to_update))

    base_dtype = x_start.dtype.base_dtype
    iter_, x_update_diff_norm_sq, x_update, _ = tf.while_loop(
        cond=_loop_cond,
        body=_loop_body,
        loop_vars=[
            tf.zeros([], dtype=np.int32, name='iter'),
            tf.zeros(
                batch_shape, dtype=base_dtype, name='x_update_diff_norm_sq'),
            tf.zeros(update_shape, dtype=base_dtype, name='x_update'),
            tf.zeros(
                update_shape, dtype=base_dtype, name='hess_matmul_x_update'),
        ])

    # Convert back x_update to the shape of x_start by transposing the leftmost
    # dimension to the rightmost.
    n = tf.rank(x_update)
    perm = tf.roll(tf.range(n), shift=-1, axis=0)
    x_update = tf.transpose(a=x_update, perm=perm)

    converged = tf.reduce_all(input_tensor=x_update_diff_norm_sq <
                              x_update_diff_norm_sq_convergence_threshold)
    return x_start + x_update, converged, iter_ / dims

def minimize(grad_and_hessian_loss_fn,
             x_start,
             tolerance,
             l1_regularizer,
             l2_regularizer=None,
             maximum_iterations=1,
             maximum_full_sweeps_per_iteration=1,
             learning_rate=None,
             name=None):
  """Minimize using Hessian-informed proximal gradient descent.

  This function solves the regularized minimization problem

  ```none
  argmin{ Loss(x)
            + l1_regularizer * ||x||_1
            + l2_regularizer * ||x||_2**2
          : x in R^n }
  ```

  where `Loss` is a convex C^2 function (typically, `Loss` is the negative log
  likelihood of a model and `x` is a vector of model coefficients).  The `Loss`
  function does not need to be supplied directly, but this optimizer does need a
  way to compute the gradient and Hessian of the Loss function at a given value
  of `x`.  The gradient and Hessian are often computationally expensive, and
  this optimizer calls them relatively few times compared with other algorithms.

  Args:
    grad_and_hessian_loss_fn: callable that takes as input a (batch of) `Tensor`
      of the same shape and dtype as `x_start` and returns the triple
      `(gradient_unregularized_loss, hessian_unregularized_loss_outer,
      hessian_unregularized_loss_middle)` as defined in the argument spec of
      `minimize_one_step`.
    x_start: (Batch of) vector-shaped, `float` `Tensor` representing the initial
      value of the argument to the `Loss` function.
    tolerance: scalar, `float` `Tensor` representing the tolerance for each
      optimization step; see the `tolerance` argument of
      `minimize_one_step`.
    l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
      regularization term (see equation above).
    l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
      regularization term (see equation above).
      Default value: `None` (i.e., no L2 regularization).
    maximum_iterations: Python integer specifying the maximum number of
      iterations of the outer loop of the optimizer.  After this many iterations
      of the outer loop, the algorithm will terminate even if the return value
      `optimal_x` has not converged.
      Default value: `1`.
    maximum_full_sweeps_per_iteration: Python integer specifying the maximum
      number of sweeps allowed in each iteration of the outer loop of the
      optimizer.  Passed as the `maximum_full_sweeps` argument to
      `minimize_one_step`.
      Default value: `1`.
    learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
      used to dampen the proximal gradient descent steps.
      Default value: `None` (i.e., factor is conceptually `1`).
    name: Python string representing the name of the TensorFlow operation.
      The default name is `"minimize"`.

  Returns:
    x: `Tensor` of the same shape and dtype as `x_start`, representing the
      (batches of) computed values of `x` which minimizes `Loss(x)`.
    is_converged: scalar, `bool` `Tensor` indicating whether the minimization
      procedure converged within the specified number of iterations across all
      batches.  Here convergence means that an iteration of the inner loop
      (`minimize_one_step`) returns `True` for its `is_converged` output value.
    iter: scalar, `int` `Tensor` indicating the actual number of iterations of
      the outer loop of the optimizer completed (i.e., number of calls to
      `minimize_one_step` before achieving convergence).

  #### References

  [1]: Jerome Friedman, Trevor Hastie and Rob Tibshirani. Regularization Paths
       for Generalized Linear Models via Coordinate Descent. _Journal of
       Statistical Software_, 33(1), 2010.
       https://www.jstatsoft.org/article/view/v033i01/v33i01.pdf

  [2]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for
       L1-regularized Logistic Regression. _Journal of Machine Learning
       Research_, 13, 2012.
       http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf
  """
  graph_deps = [
      x_start,
      l1_regularizer,
      l2_regularizer,
      maximum_iterations,
      maximum_full_sweeps_per_iteration,
      tolerance,
      learning_rate,
  ],
  with tf.compat.v1.name_scope(name, 'minimize', graph_deps):

    def _loop_cond(x_start, converged, iter_):
      del x_start
      return tf.logical_and(iter_ < maximum_iterations,
                            tf.logical_not(converged))

    def _loop_body(x_start, converged, iter_):  # pylint: disable=missing-docstring
      g, h_outer, h_middle = grad_and_hessian_loss_fn(x_start)
      x_start, converged, _ = minimize_one_step(
          gradient_unregularized_loss=g,
          hessian_unregularized_loss_outer=h_outer,
          hessian_unregularized_loss_middle=h_middle,
          x_start=x_start,
          l1_regularizer=l1_regularizer,
          l2_regularizer=l2_regularizer,
          maximum_full_sweeps=maximum_full_sweeps_per_iteration,
          tolerance=tolerance,
          learning_rate=learning_rate)
      return x_start, converged, iter_ + 1

    return tf.while_loop(
        cond=_loop_cond,
        body=_loop_body,
        loop_vars=[
            x_start,
            tf.zeros([], np.bool, name='converged'),
            tf.zeros([], np.int32, name='iter'),
        ])


# 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 vector quantized-variational autoencoder (VQ-VAE) on MNIST.

The VQ-VAE is similar to a variational autoencoder (VAE), but the latent
code Z goes through a discrete bottleneck before being passed to the encoder.
The bottleneck uses vector quantization to match the latent code to its nearest
neighbor in a codebook. To train, we minimize the weighted sum of the
reconstruction loss and a commitment loss that ensures the encoder commits to
entries in the codebook. In addition, we use exponential moving averaging (EMA)
to update the codebook for each minibatch.

#### References

[1]: Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu. Neural Discrete
     Representation Learning. In _Conference on Neural Information Processing
     Systems_, 2017. https://arxiv.org/abs/1711.00937
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import os
import time

# Dependency imports
from absl import flags
from matplotlib import cm
from matplotlib import figure
from matplotlib.backends import backend_agg
import numpy as np
from six.moves import urllib
import tensorflow as tf

from tensorflow_probability import distributions as tfd
from tensorflow.contrib.learn.python.learn.datasets import mnist
from tensorflow.python.training import moving_averages

IMAGE_SHAPE = [28, 28, 1]

flags.DEFINE_float("learning_rate",
                   default=0.001,
                   help="Initial learning rate.")
flags.DEFINE_integer("max_steps",
                     default=10000,
                     help="Number of training steps to run.")
flags.DEFINE_integer("latent_size",
                     default=1,
                     help="Number of latent variables.")
flags.DEFINE_integer("num_codes",
                     default=64,
                     help="Number of discrete codes in codebook.")
flags.DEFINE_integer("code_size",
                     default=16,
                     help="Dimension of each entry in codebook.")
flags.DEFINE_integer("base_depth",
                     default=32,
                     help="Base depth for encoder and decoder CNNs.")
flags.DEFINE_string("activation",
                    default="elu",
                    help="Activation function for all hidden layers.")
flags.DEFINE_float("beta",
                   default=0.25,
                   help="Scaling for commitment loss.")
flags.DEFINE_float("decay",
                   default=0.99,
                   help="Decay for exponential moving average.")
flags.DEFINE_integer("batch_size",
                     default=128,
                     help="Batch size.")
flags.DEFINE_string("mnist_type",
                    default="threshold",
                    help="""Type of MNIST used. Choices include 'fake_data',
                    'bernoulli' for Hugo Larochelle's randomly binarized MNIST,
                     and 'threshold' for binarized MNIST at 0.5 threshold.""")
flags.DEFINE_string("data_dir",
                    default=os.path.join(os.getenv("TEST_TMPDIR", "/tmp"),
                                         "vq_vae/data"),
                    help="Directory where data is stored (if using real data).")
flags.DEFINE_string(
    "model_dir",
    default=os.path.join(os.getenv("TEST_TMPDIR", "/tmp"), "vq_vae/"),
    help="Directory to put the model's fit.")
flags.DEFINE_integer("viz_steps",
                     default=500,
                     help="Frequency at which to save visualizations.")

FLAGS = flags.FLAGS
BERNOULLI_PATH = "http://www.cs.toronto.edu/~larocheh/public/datasets/binarized_mnist/"
FILE_TEMPLATE = "binarized_mnist_{split}.amat"

class MnistType(object):
  """MNIST types for input data."""
  FAKE_DATA = "fake_data"
  THRESHOLD = "threshold"
  BERNOULLI = "bernoulli"

class VectorQuantizer(object):
  """Creates a vector-quantizer.

  It quantizes a continuous vector under a codebook. The codebook is also known
  as "embeddings" or "memory", and it is learned using an exponential moving
  average.
  """

  def __init__(self, num_codes, code_size):
    self.num_codes = num_codes
    self.code_size = code_size
    self.codebook = tf.compat.v1.get_variable(
        "codebook",
        [num_codes, code_size],
        dtype=tf.float32,
    )
    self.ema_count = tf.compat.v1.get_variable(
        name="ema_count",
        shape=[num_codes],
        initializer=tf.compat.v1.initializers.constant(0),
        trainable=False)
    self.ema_means = tf.compat.v1.get_variable(
        name="ema_means",
        initializer=self.codebook.initialized_value(),
        trainable=False)

  def __call__(self, codes):
    """Uses codebook to find nearest neighbor for each code.

    Args:
      codes: A `float`-like `Tensor` containing the latent
        vectors to be compared to the codebook. These are rank-3 with shape
        `[batch_size, latent_size, code_size]`.

    Returns:
      nearest_codebook_entries: The 1-nearest neighbor in Euclidean distance for
        each code in the batch.
      one_hot_assignments: The one-hot vectors corresponding to the matched
        codebook entry for each code in the batch.
    """
    distances = tf.norm(
        tensor=tf.expand_dims(codes, 2) -
        tf.reshape(self.codebook, [1, 1, self.num_codes, self.code_size]),
        axis=3)
    assignments = tf.argmin(input=distances, axis=2)
    one_hot_assignments = tf.one_hot(assignments, depth=self.num_codes)
    nearest_codebook_entries = tf.reduce_sum(
        input_tensor=tf.expand_dims(one_hot_assignments, -1) *
        tf.reshape(self.codebook, [1, 1, self.num_codes, self.code_size]),
        axis=2)
    return nearest_codebook_entries, one_hot_assignments

def make_encoder(base_depth, activation, latent_size, code_size):
  """Creates the encoder function.

  Args:
    base_depth: Layer base depth in encoder net.
    activation: Activation function in hidden layers.
    latent_size: The number of latent variables in the code.
    code_size: The dimensionality of each latent variable.

  Returns:
    encoder: A `callable` mapping a `Tensor` of images to a `Tensor` of shape
      `[..., latent_size, code_size]`.
  """
  conv = functools.partial(
      tf.keras.layers.Conv2D, padding="SAME", activation=activation)

  encoder_net = tf.keras.Sequential([
      conv(base_depth, 5, 1),
      conv(base_depth, 5, 2),
      conv(2 * base_depth, 5, 1),
      conv(2 * base_depth, 5, 2),
      conv(4 * latent_size, 7, padding="VALID"),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(latent_size * code_size, activation=None),
      tf.keras.layers.Reshape([latent_size, code_size])
  ])

  def encoder(images):
    """Encodes a batch of images.

    Args:
      images: A `Tensor` representing the inputs to be encoded, of shape `[...,
        channels]`.

    Returns:
      codes: A `float`-like `Tensor` of shape `[..., latent_size, code_size]`.
        It represents latent vectors to be matched with the codebook.
    """
    images = 2 * tf.cast(images, dtype=tf.float32) - 1
    codes = encoder_net(images)
    return codes

  return encoder

def make_decoder(base_depth, activation, input_size, output_shape):
  """Creates the decoder function.

  Args:
    base_depth: Layer base depth in decoder net.
    activation: Activation function in hidden layers.
    input_size: The flattened latent input shape as an int.
    output_shape: The output image shape as a list.

  Returns:
    decoder: A `callable` mapping a `Tensor` of encodings to a
      `tfd.Distribution` instance over images.
  """
  deconv = functools.partial(
      tf.keras.layers.Conv2DTranspose, padding="SAME", activation=activation)
  conv = functools.partial(
      tf.keras.layers.Conv2D, padding="SAME", activation=activation)
  decoder_net = tf.keras.Sequential([
      tf.keras.layers.Reshape((1, 1, input_size)),
      deconv(2 * base_depth, 7, padding="VALID"),
      deconv(2 * base_depth, 5),
      deconv(2 * base_depth, 5, 2),
      deconv(base_depth, 5),
      deconv(base_depth, 5, 2),
      deconv(base_depth, 5),
      conv(output_shape[-1], 5, activation=None),
      tf.keras.layers.Reshape(output_shape),
  ])

  def decoder(codes):
    """Builds a distribution over images given codes.

    Args:
      codes: A `Tensor` representing the inputs to be decoded, of shape `[...,
        code_size]`.

    Returns:
      decoder_distribution: A multivariate `Bernoulli` distribution.
    """
    logits = decoder_net(codes)
    return tfd.Independent(tfd.Bernoulli(logits=logits),
                           reinterpreted_batch_ndims=len(output_shape),
                           name="decoder_distribution")

  return decoder

def add_ema_control_dependencies(vector_quantizer,
                                 one_hot_assignments,
                                 codes,
                                 commitment_loss,
                                 decay):
  """Add control dependencies to the commmitment loss to update the codebook.

  Args:
    vector_quantizer: An instance of the VectorQuantizer class.
    one_hot_assignments: The one-hot vectors corresponding to the matched
      codebook entry for each code in the batch.
    codes: A `float`-like `Tensor` containing the latent vectors to be compared
      to the codebook.
    commitment_loss: The commitment loss from comparing the encoder outputs to
      their neighboring codebook entries.
    decay: Decay factor for exponential moving average.

  Returns:
    commitment_loss: Commitment loss with control dependencies.
  """
  # Use an exponential moving average to update the codebook.
  updated_ema_count = moving_averages.assign_moving_average(
      vector_quantizer.ema_count,
      tf.reduce_sum(input_tensor=one_hot_assignments, axis=[0, 1]),
      decay,
      zero_debias=False)
  updated_ema_means = moving_averages.assign_moving_average(
      vector_quantizer.ema_means,
      tf.reduce_sum(
          input_tensor=tf.expand_dims(codes, 2) *
          tf.expand_dims(one_hot_assignments, 3),
          axis=[0, 1]),
      decay,
      zero_debias=False)

  # Add small value to avoid dividing by zero.
  perturbed_ema_count = updated_ema_count + 1e-5
  with tf.control_dependencies([commitment_loss]):
    update_means = tf.compat.v1.assign(
        vector_quantizer.codebook,
        updated_ema_means / perturbed_ema_count[..., tf.newaxis])
    with tf.control_dependencies([update_means]):
      return tf.identity(commitment_loss)

def save_imgs(x, fname):
  """Helper method to save a grid of images to a PNG file.

  Args:
    x: A numpy array of shape [n_images, height, width].
    fname: The filename to write to (including extension).
  """
  n = x.shape[0]
  fig = figure.Figure(figsize=(n, 1), frameon=False)
  canvas = backend_agg.FigureCanvasAgg(fig)
  for i in range(n):
    ax = fig.add_subplot(1, n, i+1)
    ax.imshow(x[i].squeeze(),
              interpolation="none",
              cmap=cm.get_cmap("binary"))
    ax.axis("off")
  canvas.print_figure(fname, format="png")
  print("saved %s" % fname)

def visualize_training(images_val,
                       reconstructed_images_val,
                       random_images_val,
                       log_dir, prefix, viz_n=10):
  """Helper method to save images visualizing model reconstructions.

  Args:
    images_val: Numpy array containing a batch of input images.
    reconstructed_images_val: Numpy array giving the expected output
      (mean) of the decoder.
    random_images_val: Optionally, a Numpy array giving the expected output
      (mean) of decoding samples from the prior, or `None`.
    log_dir: The directory to write images (Python `str`).
    prefix: A specific label for the saved visualizations, which
      determines their filenames (Python `str`).
    viz_n: The number of images from each batch to visualize (Python `int`).
  """
  save_imgs(images_val[:viz_n],
            os.path.join(log_dir, "{}_inputs.png".format(prefix)))
  save_imgs(reconstructed_images_val[:viz_n],
            os.path.join(log_dir,
                         "{}_reconstructions.png".format(prefix)))

  if random_images_val is not None:
    save_imgs(random_images_val[:viz_n],
              os.path.join(log_dir,
                           "{}_prior_samples.png".format(prefix)))

def build_fake_data(num_examples=10):
  """Builds fake MNIST-style data for unit testing."""

  class Dummy(object):
    pass

  num_examples = 10
  mnist_data = Dummy()
  mnist_data.train = Dummy()
  mnist_data.train.images = np.float32(np.random.randn(
      num_examples, np.prod(IMAGE_SHAPE)))
  mnist_data.train.labels = np.int32(np.random.permutation(
      np.arange(num_examples)))
  mnist_data.train.num_examples = num_examples
  mnist_data.validation = Dummy()
  mnist_data.validation.images = np.float32(np.random.randn(
      num_examples, np.prod(IMAGE_SHAPE)))
  mnist_data.validation.labels = np.int32(np.random.permutation(
      np.arange(num_examples)))
  mnist_data.validation.num_examples = num_examples
  return mnist_data

def download(directory, filename):
  """Downloads a file."""
  filepath = os.path.join(directory, filename)
  if tf.io.gfile.exists(filepath):
    return filepath
  if not tf.io.gfile.exists(directory):
    tf.io.gfile.makedirs(directory)
  url = os.path.join(BERNOULLI_PATH, filename)
  print("Downloading %s to %s" % (url, filepath))
  urllib.request.urlretrieve(url, filepath)
  return filepath

def load_bernoulli_mnist_dataset(directory, split_name):
  """Returns Hugo Larochelle's binary static MNIST tf.data.Dataset."""
  amat_file = download(directory, FILE_TEMPLATE.format(split=split_name))
  dataset = tf.data.TextLineDataset(amat_file)
  str_to_arr = lambda string: np.array([c == b"1" for c in string.split()])

  def _parser(s):
    booltensor = tf.compat.v1.py_func(str_to_arr, [s], tf.bool)
    reshaped = tf.reshape(booltensor, [28, 28, 1])
    return tf.cast(reshaped, dtype=tf.float32), tf.constant(0, tf.int32)

  return dataset.map(_parser)

def build_input_pipeline(data_dir, batch_size, heldout_size, mnist_type):
  """Builds an Iterator switching between train and heldout data."""
  # Build an iterator over training batches.
  if mnist_type in [MnistType.FAKE_DATA, MnistType.THRESHOLD]:
    if mnist_type == MnistType.FAKE_DATA:
      mnist_data = build_fake_data()
    else:
      mnist_data = mnist.read_data_sets(data_dir)
    training_dataset = tf.data.Dataset.from_tensor_slices(
        (mnist_data.train.images, np.int32(mnist_data.train.labels)))
    heldout_dataset = tf.data.Dataset.from_tensor_slices(
        (mnist_data.validation.images,
         np.int32(mnist_data.validation.labels)))
  elif mnist_type == MnistType.BERNOULLI:
    training_dataset = load_bernoulli_mnist_dataset(data_dir, "train")
    heldout_dataset = load_bernoulli_mnist_dataset(data_dir, "valid")
  else:
    raise ValueError("Unknown MNIST type.")

  training_batches = training_dataset.repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)

  # Build a iterator over the heldout set with batch_size=heldout_size,
  # i.e., return the entire heldout set as a constant.
  heldout_frozen = (heldout_dataset.take(heldout_size).
                    repeat().batch(heldout_size))
  heldout_iterator = tf.compat.v1.data.make_one_shot_iterator(heldout_frozen)

  # Combine these into a feedable iterator that can switch between training
  # and validation inputs.
  handle = tf.compat.v1.placeholder(tf.string, shape=[])
  feedable_iterator = tf.compat.v1.data.Iterator.from_string_handle(
      handle, training_batches.output_types, training_batches.output_shapes)
  images, labels = feedable_iterator.get_next()
  # Reshape as a pixel image and binarize pixels.
  images = tf.reshape(images, shape=[-1] + IMAGE_SHAPE)
  if mnist_type in [MnistType.FAKE_DATA, MnistType.THRESHOLD]:
    images = tf.cast(images > 0.5, dtype=tf.int32)

  return images, labels, handle, training_iterator, heldout_iterator

def main(argv):
  del argv  # unused
  FLAGS.activation = getattr(tf.nn, FLAGS.activation)
  if tf.io.gfile.exists(FLAGS.model_dir):
    tf.compat.v1.logging.warn("Deleting old log directory at {}".format(
        FLAGS.model_dir))
    tf.io.gfile.rmtree(FLAGS.model_dir)
  tf.io.gfile.makedirs(FLAGS.model_dir)

  with tf.Graph().as_default():
    # TODO(b/113163167): Speed up and tune hyperparameters for Bernoulli MNIST.
    (images, _, handle,
     training_iterator, heldout_iterator) = build_input_pipeline(
         FLAGS.data_dir, FLAGS.batch_size, heldout_size=10000,
         mnist_type=FLAGS.mnist_type)

    encoder = make_encoder(FLAGS.base_depth,
                           FLAGS.activation,
                           FLAGS.latent_size,
                           FLAGS.code_size)
    decoder = make_decoder(FLAGS.base_depth,
                           FLAGS.activation,
                           FLAGS.latent_size * FLAGS.code_size,
                           IMAGE_SHAPE)
    vector_quantizer = VectorQuantizer(FLAGS.num_codes, FLAGS.code_size)

    codes = encoder(images)
    nearest_codebook_entries, one_hot_assignments = vector_quantizer(codes)
    codes_straight_through = codes + tf.stop_gradient(
        nearest_codebook_entries - codes)
    decoder_distribution = decoder(codes_straight_through)
    reconstructed_images = decoder_distribution.mean()

    reconstruction_loss = -tf.reduce_mean(
        input_tensor=decoder_distribution.log_prob(images))
    commitment_loss = tf.reduce_mean(
        input_tensor=tf.square(codes -
                               tf.stop_gradient(nearest_codebook_entries)))
    commitment_loss = add_ema_control_dependencies(
        vector_quantizer,
        one_hot_assignments,
        codes,
        commitment_loss,
        FLAGS.decay)
    prior_dist = tfd.Multinomial(
        total_count=1.0, logits=tf.zeros([FLAGS.latent_size, FLAGS.num_codes]))
    prior_loss = -tf.reduce_mean(
        input_tensor=tf.reduce_sum(
            input_tensor=prior_dist.log_prob(one_hot_assignments), axis=1))

    loss = reconstruction_loss + FLAGS.beta * commitment_loss + prior_loss
    # Upper bound marginal negative log-likelihood as prior loss +
    # reconstruction loss.
    marginal_nll = prior_loss + reconstruction_loss

    tf.compat.v1.summary.scalar("losses/total_loss", loss)
    tf.compat.v1.summary.scalar("losses/reconstruction_loss",
                                reconstruction_loss)
    tf.compat.v1.summary.scalar("losses/prior_loss", prior_loss)
    tf.compat.v1.summary.scalar("losses/commitment_loss",
                                FLAGS.beta * commitment_loss)

    # Decode samples from a uniform prior for visualization.
    prior_samples = tf.reduce_sum(
        input_tensor=tf.expand_dims(prior_dist.sample(10), -1) *
        tf.reshape(vector_quantizer.codebook,
                   [1, 1, FLAGS.num_codes, FLAGS.code_size]),
        axis=2)
    decoded_distribution_given_random_prior = decoder(prior_samples)
    random_images = decoded_distribution_given_random_prior.mean()

    # Perform inference by minimizing the loss function.
    optimizer = tf.compat.v1.train.AdamOptimizer(FLAGS.learning_rate)
    train_op = optimizer.minimize(loss)

    summary = tf.compat.v1.summary.merge_all()
    init = tf.compat.v1.global_variables_initializer()
    saver = tf.compat.v1.train.Saver()
    with tf.compat.v1.Session() as sess:
      summary_writer = tf.compat.v1.summary.FileWriter(FLAGS.model_dir,
                                                       sess.graph)
      sess.run(init)

      # Run the training loop.
      train_handle = sess.run(training_iterator.string_handle())
      heldout_handle = sess.run(heldout_iterator.string_handle())
      for step in range(FLAGS.max_steps):
        start_time = time.time()
        _, loss_value = sess.run([train_op, loss],
                                 feed_dict={handle: train_handle})
        duration = time.time() - start_time
        if step % 100 == 0:
          marginal_nll_val = sess.run(marginal_nll,
                                      feed_dict={handle: heldout_handle})
          print("Step: {:>3d} Training Loss: {:.3f} Heldout NLL: {:.3f} "
                "({:.3f} sec)".format(step, loss_value, marginal_nll_val,
                                      duration))

          # Update the events file.
          summary_str = sess.run(summary, feed_dict={handle: train_handle})
          summary_writer.add_summary(summary_str, step)
          summary_writer.flush()

        # Periodically save a checkpoint and visualize model progress.
        if (step + 1) % FLAGS.viz_steps == 0 or (step + 1) == FLAGS.max_steps:
          checkpoint_file = os.path.join(FLAGS.model_dir, "model.ckpt")
          saver.save(sess, checkpoint_file, global_step=step)

          # Visualize inputs and model reconstructions from the training set.
          images_val, reconstructions_val, random_images_val = sess.run(
              (images, reconstructed_images, random_images),
              feed_dict={handle: train_handle})
          visualize_training(images_val,
                             reconstructions_val,
                             random_images_val,
                             log_dir=FLAGS.model_dir,
                             prefix="step{:05d}_train".format(step))

          # Visualize inputs and model reconstructions from the validation set.
          heldout_images_val, heldout_reconstructions_val = sess.run(
              (images, reconstructed_images),
              feed_dict={handle: heldout_handle})
          visualize_training(heldout_images_val,
                             heldout_reconstructions_val,
                             None,
                             log_dir=FLAGS.model_dir,
                             prefix="step{:05d}_validation".format(step))

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.
# ============================================================================
"""Utility functions for dtypes."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import numpy as np

import tensorflow.compat.v2 as tf

__all__ = [
    'as_numpy_dtype',
    'assert_same_float_dtype',
    'base_dtype',
    'base_equal',
    'common_dtype',
    'is_bool',
    'is_complex',
    'is_floating',
    'is_integer',
    'max',
    'min',
    'name',
    'size',
]

def as_numpy_dtype(dtype):
  """Returns a `np.dtype` based on this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'as_numpy_dtype'):
    return dtype.as_numpy_dtype
  return dtype

def base_dtype(dtype):
  """Returns a non-reference `dtype` based on this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'base_dtype'):
    return dtype.base_dtype
  return dtype

def base_equal(a, b):
  """Returns `True` if base dtypes are identical."""
  return base_dtype(a) == base_dtype(b)

def common_dtype(args_list, preferred_dtype=None):
  """Returns explict dtype from `args_list` if there is one."""
  dtype = None
  for a in tf.nest.flatten(args_list):
    if hasattr(a, 'dtype'):
      dt = as_numpy_dtype(a.dtype)
    else:
      continue
    if dtype is None:
      dtype = dt
    elif dtype != dt:
      raise TypeError('Found incompatible dtypes, {} and {}.'.format(dtype, dt))
  return preferred_dtype if dtype is None else tf.as_dtype(dtype)

def is_bool(dtype):
  """Returns whether this is a boolean data type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_bool'):
    return dtype.is_bool
  # We use `kind` because:
  # np.issubdtype(np.uint8, np.bool) == True.
  return np.dtype(dtype).kind == 'b'

def is_complex(dtype):
  """Returns whether this is a complex floating point type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_complex'):
    return dtype.is_complex
  return np.issubdtype(np.dtype(dtype), np.complex)

def is_floating(dtype):
  """Returns whether this is a (non-quantized, real) floating point type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_floating'):
    return dtype.is_floating
  return np.issubdtype(np.dtype(dtype), np.float)

def is_integer(dtype):
  """Returns whether this is a (non-quantized) integer type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_integer'):
    return dtype.is_integer
  return np.issubdtype(np.dtype(dtype), np.integer)

def max(dtype):  # pylint: disable=redefined-builtin
  """Returns the maximum representable value in this data type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'max'):
    return dtype.max
  use_finfo = is_floating(dtype) or is_complex(dtype)
  return np.finfo(dtype).max if use_finfo else np.iinfo(dtype).max

def min(dtype):  # pylint: disable=redefined-builtin
  """Returns the minimum representable value in this data type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'min'):
    return dtype.min
  use_finfo = is_floating(dtype) or is_complex(dtype)
  return np.finfo(dtype).min if use_finfo else np.iinfo(dtype).min

def name(dtype):
  """Returns the string name for this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'name'):
    return dtype.name
  if hasattr(dtype, '__name__'):
    return dtype.__name__
  return str(dtype)

def size(dtype):
  """Returns the number of bytes to represent this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'size'):
    return dtype.size
  return np.dtype(dtype).itemsize

def _assert_same_base_type(items, expected_type=None):
  r"""Asserts all items are of the same base type.

  Args:
    items: List of graph items (e.g., `Variable`, `Tensor`, `SparseTensor`,
        `Operation`, or `IndexedSlices`). Can include `None` elements, which
        will be ignored.
    expected_type: Expected type. If not specified, assert all items are
        of the same base type.

  Returns:
    Validated type, or none if neither expected_type nor items provided.

  Raises:
    ValueError: If any types do not match.
  """
  original_expected_type = expected_type
  mismatch = False
  for item in items:
    if item is not None:
      item_type = base_dtype(item.dtype)
      if not expected_type:
        expected_type = item_type
      elif expected_type != item_type:
        mismatch = True
        break
  if mismatch:
    # Loop back through and build up an informative error message (this is very
    # slow, so we don't do it unless we found an error above).
    expected_type = original_expected_type
    original_item_str = None
    get_name = lambda x: x.name if hasattr(x, 'name') else str(x)
    for item in items:
      if item is not None:
        item_type = base_dtype(item.dtype)
        if not expected_type:
          expected_type = item_type
          original_item_str = get_name(item)
        elif expected_type != item_type:
          raise ValueError(
              '{}, type={}, must be of the same type ({}){}.'.format(
                  get_name(item),
                  item_type,
                  expected_type,
                  ((' as {}'.format(original_item_str))
                   if original_item_str else '')))
    return expected_type  # Should be unreachable
  else:
    return expected_type

def assert_same_float_dtype(tensors=None, dtype=None):
  """Validate and return float type based on `tensors` and `dtype`.

  For ops such as matrix multiplication, inputs and weights must be of the
  same float type. This function validates that all `tensors` are the same type,
  validates that type is `dtype` (if supplied), and returns the type. Type must
  be a floating point type. If neither `tensors` nor `dtype` is supplied,
  the function will return `dtypes.float32`.

  Args:
    tensors: Tensors of input values. Can include `None` elements, which will
      be ignored.
    dtype: Expected type.

  Returns:
    Validated type.

  Raises:
    ValueError: if neither `tensors` nor `dtype` is supplied, or result is not
      float, or the common type of the inputs is not a floating point type.
  """
  if tensors:
    dtype = _assert_same_base_type(tensors, dtype)
  if not dtype:
    dtype = tf.float32
  elif not is_floating(dtype):
    raise ValueError('Expected floating point type, got {}.'.format(dtype))
  return dtype


# 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 OneHotCategorical distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

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 distribution_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

class OneHotCategorical(distribution.Distribution):
  """OneHotCategorical distribution.

  The categorical distribution is parameterized by the log-probabilities
  of a set of classes. The difference between OneHotCategorical and Categorical
  distributions is that OneHotCategorical is a discrete distribution over
  one-hot bit vectors whereas Categorical is a discrete distribution over
  positive integers. OneHotCategorical is equivalent to Categorical except
  Categorical has event_dim=() while OneHotCategorical has event_dim=K, where
  K is the number of classes.

  This class provides methods to create indexed batches of OneHotCategorical
  distributions. If the provided `logits` or `probs` is rank 2 or higher, for
  every fixed set of leading dimensions, the last dimension represents one
  single OneHotCategorical distribution. When calling distribution
  functions (e.g. `dist.prob(x)`), `logits` and `x` are broadcast to the
  same shape (if possible). In all cases, the last dimension of `logits,x`
  represents single OneHotCategorical distributions.

  #### Examples

  Creates a 3-class distribution, with the 2nd class, the most likely to be
  drawn from.

  ```python
  p = [0.1, 0.5, 0.4]
  dist = OneHotCategorical(probs=p)
  ```

  Creates a 3-class distribution, with the 2nd class the most likely to be
  drawn from, using logits.

  ```python
  logits = [-2, 2, 0]
  dist = OneHotCategorical(logits=logits)
  ```

  Creates a 3-class distribution, with the 3rd class is most likely to be drawn.

  ```python
  # counts is a scalar.
  p = [0.1, 0.4, 0.5]
  dist = OneHotCategorical(probs=p)
  dist.prob([0,1,0])  # Shape []

  # p will be broadcast to [[0.1, 0.4, 0.5], [0.1, 0.4, 0.5]] to match.
  samples = [[0,1,0], [1,0,0]]
  dist.prob(samples)  # Shape [2]
  ```

  """

  def __init__(self,
               logits=None,
               probs=None,
               dtype=tf.int32,
               validate_args=False,
               allow_nan_stats=True,
               name="OneHotCategorical"):
    """Initialize OneHotCategorical distributions using class log-probabilities.

    Args:
      logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities of a
        set of Categorical distributions. The first `N - 1` dimensions index
        into a batch of independent distributions and the last dimension
        represents a vector of logits for each class. Only one of `logits` or
        `probs` should be passed in.
      probs: An N-D `Tensor`, `N >= 1`, representing the probabilities of a set
        of Categorical distributions. The first `N - 1` dimensions index into a
        batch of independent distributions and the last dimension represents a
        vector of probabilities for each class. Only one of `logits` or `probs`
        should be passed in.
      dtype: The type of the event samples (default: int32).
      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())
    with tf.name_scope(name) as name:
      self._logits, self._probs = distribution_util.get_logits_and_probs(
          name=name, logits=logits, probs=probs, validate_args=validate_args,
          multidimensional=True)

      logits_shape_static = tensorshape_util.with_rank_at_least(
          self._logits.shape, 1)
      if tensorshape_util.rank(logits_shape_static) is not None:
        self._batch_rank = tf.convert_to_tensor(
            value=tensorshape_util.rank(logits_shape_static) - 1,
            dtype=tf.int32,
            name="batch_rank")
      else:
        with tf.name_scope("batch_rank"):
          self._batch_rank = tf.rank(self._logits) - 1

      with tf.name_scope("event_size"):
        self._event_size = tf.shape(input=self._logits)[-1]

    super(OneHotCategorical, self).__init__(
        dtype=dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._logits, self._probs],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(logits=1, probs=1)

  @property
  def event_size(self):
    """Scalar `int32` tensor: the number of classes."""
    return self._event_size

  @property
  def logits(self):
    """Vector of coordinatewise logits."""
    return self._logits

  @property
  def probs(self):
    """Vector of coordinatewise probabilities."""
    return self._probs

  def _batch_shape_tensor(self):
    return tf.shape(input=self.logits)[:-1]

  def _batch_shape(self):
    return self.logits.shape[:-1]

  def _event_shape_tensor(self):
    return tf.shape(input=self.logits)[-1:]

  def _event_shape(self):
    return tensorshape_util.with_rank_at_least(self.logits.shape, 1)[-1:]

  def _sample_n(self, n, seed=None):
    sample_shape = tf.concat([[n], tf.shape(input=self.logits)], 0)
    logits = self.logits
    if tensorshape_util.rank(logits.shape) == 2:
      logits_2d = logits
    else:
      logits_2d = tf.reshape(logits, [-1, self.event_size])
    samples = tf.random.categorical(logits_2d, n, seed=seed)
    samples = tf.transpose(a=samples)
    samples = tf.one_hot(samples, self.event_size, dtype=self.dtype)
    ret = tf.reshape(samples, sample_shape)
    return ret

  def _log_prob(self, x):
    x = tf.cast(x, self.logits.dtype)
    x = self._assert_valid_sample(x)
    # broadcast logits or x if need be.
    logits = self.logits
    if (not tensorshape_util.is_fully_defined(x.shape) or
        not tensorshape_util.is_fully_defined(logits.shape) or
        x.shape != logits.shape):
      logits = tf.ones_like(x, dtype=logits.dtype) * logits
      x = tf.ones_like(logits, dtype=x.dtype) * x

    logits_shape = tf.shape(input=tf.reduce_sum(input_tensor=logits, axis=-1))
    logits_2d = tf.reshape(logits, [-1, self.event_size])
    x_2d = tf.reshape(x, [-1, self.event_size])
    ret = -tf.nn.softmax_cross_entropy_with_logits(
        labels=tf.stop_gradient(x_2d),
        logits=logits_2d)
    # Reshape back to user-supplied batch and sample dims prior to 2D reshape.
    ret = tf.reshape(ret, logits_shape)
    return ret

  def _entropy(self):
    return -tf.reduce_sum(
        input_tensor=tf.nn.log_softmax(self.logits) * self.probs, axis=-1)

  def _mean(self):
    return self.probs

  def _mode(self):
    ret = tf.argmax(input=self.logits, axis=self._batch_rank)
    ret = tf.one_hot(ret, self.event_size, dtype=self.dtype)
    tensorshape_util.set_shape(ret, self.logits.shape)
    return ret

  def _covariance(self):
    p = self.probs
    ret = -tf.matmul(p[..., None], p[..., None, :])
    return tf.linalg.set_diag(ret, self._variance())

  def _variance(self):
    return self.probs * (1. - self.probs)

  def _assert_valid_sample(self, x):
    if not self.validate_args:
      return x
    return distribution_util.with_dependencies([
        assert_util.assert_non_positive(x),
        assert_util.assert_near(
            tf.zeros([], dtype=self.logits.dtype),
            tf.reduce_logsumexp(input_tensor=x, axis=[-1])),
    ], x)

@kullback_leibler.RegisterKL(OneHotCategorical, OneHotCategorical)
def _kl_categorical_categorical(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a, b OneHotCategorical.

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

  Returns:
    Batchwise KL(a || b)
  """
  with tf.name_scope(name or "kl_categorical_categorical"):
    # sum(p ln(p / q))
    return tf.reduce_sum(
        input_tensor=tf.nn.softmax(a.logits) *
        (tf.nn.log_softmax(a.logits) - tf.nn.log_softmax(b.logits)),
        axis=-1)