# 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.
# ============================================================================
"""Vectorized Exponential distribution class, directly using LinearOperator."""

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

import tensorflow.compat.v2 as tf
from tensorflow_probability.python.bijectors import affine_linear_operator as affine_linear_operator_bijector
from tensorflow_probability.python.distributions import exponential
from tensorflow_probability.python.distributions import transformed_distribution
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import tensorshape_util

__all__ = ["VectorExponentialLinearOperator"]

_mvn_sample_note = """
`value` is a batch vector with compatible shape if `value` is a `Tensor` whose
shape can be broadcast up to either:

```python
self.batch_shape + self.event_shape
```

or

```python
[M1, ..., Mm] + self.batch_shape + self.event_shape
```

"""

class VectorExponentialLinearOperator(
    transformed_distribution.TransformedDistribution):
  """The vectorization of the Exponential distribution on `R^k`.

  The vector exponential distribution is defined over a subset of `R^k`, and
  parameterized by a (batch of) length-`k` `loc` vector and a (batch of) `k x k`
  `scale` matrix:  `covariance = scale @ scale.T`, where `@` denotes
  matrix-multiplication.

  #### Mathematical Details

  The probability density function (pdf) is

  ```none
  pdf(y; loc, scale) = exp(-||x||_1) / Z,  for y in S(loc, scale),
  x = inv(scale) @ (y - loc),
  Z = |det(scale)|,
  ```

  where:

  * `loc` is a vector in `R^k`,
  * `scale` is a linear operator in `R^{k x k}`, `cov = scale @ scale.T`,
  * `S = {loc + scale @ x : x in R^k, x_1 > 0, ..., x_k > 0}`, is an image of
     the positive half-space,
  * `||x||_1` denotes the `l1` norm of `x`, `sum_i |x_i|`,
  * `Z` denotes the normalization constant.

  The VectorExponential distribution is a member of the [location-scale
  family](https://en.wikipedia.org/wiki/Location-scale_family), i.e., it can be
  constructed as,

  ```none
  X = (X_1, ..., X_k), each X_i ~ Exponential(rate=1)
  Y = (Y_1, ...,Y_k) = scale @ X + loc
  ```

  #### About `VectorExponential` and `Vector` distributions in TensorFlow.

  The `VectorExponential` is a non-standard distribution that has useful
  properties.

  The marginals `Y_1, ..., Y_k` are *not* Exponential random variables, due to
  the fact that the sum of Exponential random variables is not Exponential.

  Instead, `Y` is a vector whose components are linear combinations of
  Exponential random variables.  Thus, `Y` lives in the vector space generated
  by `vectors` of Exponential distributions.  This allows the user to decide the
  mean and covariance (by setting `loc` and `scale`), while preserving some
  properties of the Exponential distribution.  In particular, the tails of `Y_i`
  will be (up to polynomial factors) exponentially decaying.

  To see this last statement, note that the pdf of `Y_i` is the convolution of
  the pdf of `k` independent Exponential random variables.  One can then show by
  induction that distributions with exponential (up to polynomial factors) tails
  are closed under convolution.

  #### Examples

  ```python
  tfd = tfp.distributions

  # Initialize a single 2-variate VectorExponential, supported on
  # {(x, y) in R^2 : x > 0, y > 0}.
  mat = [[1.0, 0.1],
         [0.1, 1.0]]

  vex = tfd.VectorExponentialLinearOperator(
      scale=tf.linalg.LinearOperatorFullMatrix(mat))

  # Compute the pdf of an`R^2` observation; return a scalar.
  vex.prob([1., 2.]).eval()  # shape: []

  # Initialize a 2-batch of 3-variate Vector Exponential's.
  mu = [[1., 2, 3],
        [1., 0, 0]]              # shape: [2, 3]
  scale_diag = [[1., 2, 3],
                [0.5, 1, 1.5]]     # shape: [2, 3]

  vex = tfd.VectorExponentialLinearOperator(
      loc=mu,
      scale=tf.linalg.LinearOperatorDiag(scale_diag))

  # Compute the pdf of two `R^3` observations; return a length-2 vector.
  x = [[1.9, 2.2, 3.1],
       [10., 1.0, 9.0]]     # shape: [2, 3]
  vex.prob(x).eval()    # shape: [2]
  ```

  """

  def __init__(self,
               loc=None,
               scale=None,
               validate_args=False,
               allow_nan_stats=True,
               name="VectorExponentialLinearOperator"):
    """Construct Vector Exponential distribution supported on a subset of `R^k`.

    The `batch_shape` is the broadcast shape between `loc` and `scale`
    arguments.

    The `event_shape` is given by last dimension of the matrix implied by
    `scale`. The last dimension of `loc` (if provided) must broadcast with this.

    Recall that `covariance = scale @ scale.T`.

    Additional leading dimensions (if any) will index batches.

    Args:
      loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
        implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
        `b >= 0` and `k` is the event size.
      scale: Instance of `LinearOperator` with same `dtype` as `loc` and shape
        `[B1, ..., Bb, k, k]`.
      validate_args: Python `bool`, default `False`. Whether to validate input
        with asserts. If `validate_args` is `False`, and the inputs are
        invalid, correct behavior is not guaranteed.
      allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
        exception if a statistic (e.g. mean/mode/etc...) is undefined for any
        batch member If `True`, batch members with valid parameters leading to
        undefined statistics will return NaN for this statistic.
      name: The name to give Ops created by the initializer.

    Raises:
      ValueError: if `scale` is unspecified.
      TypeError: if not `scale.dtype.is_floating`
    """
    parameters = dict(locals())
    if scale is None:
      raise ValueError("Missing required `scale` parameter.")
    if not dtype_util.is_floating(scale.dtype):
      raise TypeError("`scale` parameter must have floating-point dtype.")

    with tf.name_scope(name) as name:
      # Since expand_dims doesn't preserve constant-ness, we obtain the
      # non-dynamic value if possible.
      loc = loc if loc is None else tf.convert_to_tensor(
          value=loc, name="loc", dtype=scale.dtype)
      batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(
          loc, scale)

      super(VectorExponentialLinearOperator, self).__init__(
          distribution=exponential.Exponential(
              rate=tf.ones([], dtype=scale.dtype),
              allow_nan_stats=allow_nan_stats),
          bijector=affine_linear_operator_bijector.AffineLinearOperator(
              shift=loc, scale=scale, validate_args=validate_args),
          batch_shape=batch_shape,
          event_shape=event_shape,
          validate_args=validate_args,
          name=name)
      self._parameters = parameters

  @property
  def loc(self):
    """The `loc` `Tensor` in `Y = scale @ X + loc`."""
    return self.bijector.shift

  @property
  def scale(self):
    """The `scale` `LinearOperator` in `Y = scale @ X + loc`."""
    return self.bijector.scale

  @distribution_util.AppendDocstring(_mvn_sample_note)
  def _log_prob(self, x):
    return super(VectorExponentialLinearOperator, self)._log_prob(x)

  @distribution_util.AppendDocstring(_mvn_sample_note)
  def _prob(self, x):
    return super(VectorExponentialLinearOperator, self)._prob(x)

  def _mean(self):
    # Let
    #   W = (w1,...,wk), with wj ~ iid Exponential(0, 1).
    # Then this distribution is
    #   X = loc + LW,
    # and then E[X] = loc + L1, where 1 is the vector of ones.
    scale_x_ones = self.bijector.scale.matvec(
        tf.ones(self._mode_mean_shape(), self.dtype))

    if self.loc is None:
      return scale_x_ones

    return tf.identity(self.loc) + scale_x_ones

  def _covariance(self):
    # Let
    #   W = (w1,...,wk), with wj ~ iid Exponential(0, 1).
    # Then this distribution is
    #   X = loc + LW,
    # and then since Cov(wi, wj) = 1 if i=j, and 0 otherwise,
    #   Cov(X) = L Cov(W W^T) L^T = L L^T.
    if distribution_util.is_diagonal_scale(self.scale):
      return tf.linalg.diag(tf.square(self.scale.diag_part()))
    else:
      return self.scale.matmul(self.scale.to_dense(), adjoint_arg=True)

  def _variance(self):
    if distribution_util.is_diagonal_scale(self.scale):
      return tf.square(self.scale.diag_part())
    elif (isinstance(self.scale, tf.linalg.LinearOperatorLowRankUpdate) and
          self.scale.is_self_adjoint):
      return tf.linalg.diag_part(self.scale.matmul(self.scale.to_dense()))
    else:
      return tf.linalg.diag_part(
          self.scale.matmul(self.scale.to_dense(), adjoint_arg=True))

  def _stddev(self):
    if distribution_util.is_diagonal_scale(self.scale):
      return tf.abs(self.scale.diag_part())
    elif (isinstance(self.scale, tf.linalg.LinearOperatorLowRankUpdate) and
          self.scale.is_self_adjoint):
      return tf.sqrt(
          tf.linalg.diag_part(self.scale.matmul(self.scale.to_dense())))
    else:
      return tf.sqrt(
          tf.linalg.diag_part(
              self.scale.matmul(self.scale.to_dense(), adjoint_arg=True)))

  def _mode(self):
    scale_x_zeros = self.bijector.scale.matvec(
        tf.zeros(self._mode_mean_shape(), self.dtype))

    if self.loc is None:
      return scale_x_zeros

    return tf.identity(self.loc) + scale_x_zeros

  def _mode_mean_shape(self):
    """Shape for the mode/mean Tensors."""
    shape = tensorshape_util.concatenate(self.batch_shape, self.event_shape)
    has_static_shape = tensorshape_util.is_fully_defined(shape)
    if not has_static_shape:
      shape = tf.concat([
          self.batch_shape_tensor(),
          self.event_shape_tensor(),
      ], 0)
    return shape


# 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.
# ============================================================================
"""Methods for forecasting in StructuralTimeSeries models."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.sts.internal import util as sts_util

def _prefer_static_event_ndims(distribution):
  if distribution.event_shape.ndims is not None:
    return distribution.event_shape.ndims
  else:
    return tf.size(input=distribution.event_shape_tensor())

def one_step_predictive(model, observed_time_series, parameter_samples):
  """Compute one-step-ahead predictive distributions for all timesteps.

  Given samples from the posterior over parameters, return the predictive
  distribution over observations at each time `T`, given observations up
  through time `T-1`.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
      mask `Tensor` to encode the locations of missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.

  Returns:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_timesteps] and
      batch shape `concat([sample_shape, model.batch_shape])`, with
      `num_posterior_draws` mixture components. The `t`th step represents the
      forecast distribution `p(observed_time_series[t] |
      observed_time_series[0:t-1], parameter_samples)`.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `one_step_predictive`, we construct a
  one-step-ahead predictive distribution:

  ```python
    one_step_predictive_dist = tfp.sts.one_step_predictive(
      model, observed_time_series, parameter_samples=samples)

    predictive_means = one_step_predictive_dist.mean()
    predictive_scales = one_step_predictive_dist.stddev()
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    one_step_predictive_dist = tfp.sts.one_step_predictive(
      model, observed_time_series, parameter_samples=samples)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_one_step_predictive(observed_time_series,
                                 forecast_mean,
                                 forecast_scale):
      plt.figure(figsize=(12, 6))
      num_timesteps = forecast_mean.shape[-1]
      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(observed_time_series, label="observed time series", color=c1)
      plt.plot(forecast_mean, label="one-step prediction", color=c2)
      plt.fill_between(np.arange(num_timesteps),
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale,
                       alpha=0.1, color=c2)
      plt.legend()

    plot_one_step_predictive(observed_time_series,
                             forecast_mean=predictive_means,
                             forecast_scale=predictive_scales)
  ```

  To detect anomalous timesteps, we check whether the observed value at each
  step is within a 95% predictive interval, i.e., two standard deviations from
  the mean:

  ```python
    z_scores = ((observed_time_series[..., 1:] - predictive_means[..., :-1])
                 / predictive_scales[..., :-1])
    anomalous_timesteps = tf.boolean_mask(
        tf.range(1, num_timesteps),
        tf.abs(z_scores) > 2.0)
  ```

  """

  with tf.compat.v1.name_scope(
      'one_step_predictive', values=[observed_time_series, parameter_samples]):

    [
        observed_time_series,
        is_missing
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run filtering over the training timesteps to extract the
    # predictive means and variances.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    lgssm = model.make_state_space_model(
        num_timesteps=num_timesteps, param_vals=parameter_samples)
    (_, _, _, _, _, observation_means, observation_covs
    ) = lgssm.forward_filter(observed_time_series, mask=is_missing)

    # Squeeze dims to convert from LGSSM's event shape `[num_timesteps, 1]`
    # to a scalar time series.
    return sts_util.mix_over_posterior_draws(
        means=observation_means[..., 0],
        variances=observation_covs[..., 0, 0])

def forecast(model,
             observed_time_series,
             parameter_samples,
             num_steps_forecast):
  """Construct predictive distribution over future observations.

  Given samples from the posterior over parameters, return the predictive
  distribution over future observations for num_steps_forecast timesteps.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
      mask `Tensor` to encode the locations of missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.
    num_steps_forecast: scalar `int` `Tensor` number of steps to forecast.

  Returns:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_steps_forecast, 1] and batch shape
      `concat([sample_shape, model.batch_shape])`, with `num_posterior_draws`
      mixture components.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `forecast`, we construct a forecast
  distribution:

  ```python
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                     parameter_samples=samples,
                                     num_steps_forecast=50)

    forecast_mean = forecast_dist.mean()[..., 0]  # shape: [50]
    forecast_scale = forecast_dist.stddev()[..., 0]  # shape: [50]
    forecast_samples = forecast_dist.sample(10)[..., 0]  # shape: [10, 50]
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                         parameter_samples=samples,
                                         num_steps_forecast=50)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_forecast(observed_time_series,
                      forecast_mean,
                      forecast_scale,
                      forecast_samples):
      plt.figure(figsize=(12, 6))

      num_steps = observed_time_series.shape[-1]
      num_steps_forecast = forecast_mean.shape[-1]
      num_steps_train = num_steps - num_steps_forecast

      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(np.arange(num_steps), observed_time_series,
               lw=2, color=c1, label='ground truth')

      forecast_steps = np.arange(num_steps_train,
                       num_steps_train+num_steps_forecast)
      plt.plot(forecast_steps, forecast_samples.T, lw=1, color=c2, alpha=0.1)
      plt.plot(forecast_steps, forecast_mean, lw=2, ls='--', color=c2,
               label='forecast')
      plt.fill_between(forecast_steps,
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale, color=c2, alpha=0.2)

      plt.xlim([0, num_steps])
      plt.legend()

    plot_forecast(observed_time_series,
                  forecast_mean=forecast_mean,
                  forecast_scale=forecast_scale,
                  forecast_samples=forecast_samples)
  ```

  """

  with tf.compat.v1.name_scope(
      'forecast',
      values=[observed_time_series, parameter_samples, num_steps_forecast]):
    [
        observed_time_series,
        mask
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run filtering over the observed timesteps to extract the
    # latent state posterior at timestep T+1 (i.e., the final
    # filtering distribution, pushed through the transition model).
    # This is the prior for the forecast model ("today's prior
    # is yesterday's posterior").
    num_observed_steps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    observed_data_ssm = model.make_state_space_model(
        num_timesteps=num_observed_steps, param_vals=parameter_samples)
    (_, _, _, predictive_means, predictive_covs, _, _
    ) = observed_data_ssm.forward_filter(observed_time_series, mask=mask)

    # Build a batch of state-space models over the forecast period. Because
    # we'll use MixtureSameFamily to mix over the posterior draws, we need to
    # do some shenanigans to move the `[num_posterior_draws]` batch dimension
    # from the leftmost to the rightmost side of the model's batch shape.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate `move_dimension` calls here.
    parameter_samples = model._canonicalize_param_vals_as_map(parameter_samples)  # pylint: disable=protected-access
    parameter_samples_with_reordered_batch_dimension = {
        param.name: dist_util.move_dimension(
            parameter_samples[param.name],
            0, -(1 + _prefer_static_event_ndims(param.prior)))
        for param in model.parameters}
    forecast_prior = tfd.MultivariateNormalFullCovariance(
        loc=dist_util.move_dimension(predictive_means[..., -1, :], 0, -2),
        covariance_matrix=dist_util.move_dimension(
            predictive_covs[..., -1, :, :], 0, -3))

    # Ugly hack: because we moved `num_posterior_draws` to the trailing (rather
    # than leading) dimension of parameters, the parameter batch shapes no
    # longer broadcast against the `constant_offset` attribute used in `sts.Sum`
    # models. We fix this by manually adding an extra broadcasting dim to
    # `constant_offset` if present.
    # The root cause of this hack is that we mucked with param dimensions above
    # and are now passing params that are 'invalid' in the sense that they don't
    # match the shapes of the model's param priors. The fix (as above) will be
    # to update MixtureSameFamily so we can avoid changing param dimensions
    # altogether.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate this hack.
    kwargs = {}
    if hasattr(model, 'constant_offset'):
      kwargs['constant_offset'] = tf.convert_to_tensor(
          value=model.constant_offset,
          dtype=forecast_prior.dtype)[..., tf.newaxis]

    # We assume that any STS model that has a `constant_offset` attribute
    # will allow it to be overridden as a kwarg. This is currently just
    # `sts.Sum`.
    # TODO(b/120245392): when kwargs hack is removed, switch back to calling
    # the public version of `_make_state_space_model`.
    forecast_ssm = model._make_state_space_model(  # pylint: disable=protected-access
        num_timesteps=num_steps_forecast,
        param_map=parameter_samples_with_reordered_batch_dimension,
        initial_state_prior=forecast_prior,
        initial_step=num_observed_steps,
        **kwargs)

    num_posterior_draws = dist_util.prefer_static_value(
        forecast_ssm.batch_shape_tensor())[-1]
    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([num_posterior_draws], dtype=forecast_ssm.dtype)),
        components_distribution=forecast_ssm)


# 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.
# ============================================================================
"""Numpy implementations of TensorFlow functions."""

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

# Dependency imports
import numpy as np

import tensorflow as tf

from tensorflow_probability.python.internal.backend.numpy.internal import utils

scipy_special = utils.try_import('scipy.special')

__all__ = [
    'abs',
    'accumulate_n',
    'acos',
    'acosh',
    'add',
    'add_n',
    'angle',
    'argmax',
    'argmin',
    'asin',
    'asinh',
    'atan',
    'atan2',
    'atanh',
    'bessel_i0',
    'bessel_i0e',
    'bessel_i1',
    'bessel_i1e',
    'betainc',
    'bincount',
    'ceil',
    # 'confusion_matrix',
    'conj',
    'cos',
    'cosh',
    'count_nonzero',
    'cumprod',
    'cumsum',
    'digamma',
    'divide',
    'divide_no_nan',
    'equal',
    'erf',
    'erfc',
    'exp',
    'expm1',
    'floor',
    'floordiv',
    'greater',
    'greater_equal',
    'igamma',
    'igammac',
    'imag',
    # 'in_top_k',
    # 'invert_permutation',
    'is_finite',
    'is_inf',
    'is_nan',
    # 'is_non_decreasing',
    # 'is_strictly_increasing',
    # 'l2_normalize',
    'lbeta',
    'less',
    'less_equal',
    'lgamma',
    'log',
    'log1p',
    'log_sigmoid',
    'log_softmax',
    'logical_and',
    'logical_not',
    'logical_or',
    'logical_xor',
    'maximum',
    'minimum',
    'multiply',
    'multiply_no_nan',
    'negative',
    # 'nextafter',
    'not_equal',
    'polygamma',
    'polyval',
    'pow',
    'real',
    'reciprocal',
    'reduce_all',
    'reduce_any',
    # 'reduce_euclidean_norm',
    'reduce_logsumexp',
    'reduce_max',
    'reduce_mean',
    'reduce_min',
    'reduce_prod',
    'reduce_std',
    'reduce_sum',
    'reduce_variance',
    'rint',
    'round',
    'rsqrt',
    # 'scalar_mul',
    # 'segment_max',
    # 'segment_mean',
    # 'segment_min',
    # 'segment_prod',
    # 'segment_sum',
    'sigmoid',
    'sign',
    'sin',
    'sinh',
    'softmax',
    'softplus',
    'softsign',
    'sqrt',
    'square',
    'squared_difference',
    'subtract',
    'tan',
    'tanh',
    # 'top_k',
    'truediv',
    # 'unsorted_segment_max',
    # 'unsorted_segment_mean',
    # 'unsorted_segment_min',
    # 'unsorted_segment_prod',
    # 'unsorted_segment_sqrt_n',
    # 'unsorted_segment_sum',
    'xdivy',
    'xlogy',
    # 'zero_fraction',
    'zeta',
]

def _astuple(x):
  try:
    return tuple(x)
  except TypeError:
    return x

def _bincount(arr, weights=None, minlength=None, maxlength=None,  # pylint: disable=unused-argument
              dtype=tf.int32, name=None):  # pylint: disable=unused-argument
  return np.bincount(arr, weights, minlength).astype(utils.numpy_dtype(dtype))

def _lbeta(x, name=None):  # pylint: disable=unused-argument
  x = np.array(x)
  return scipy_special.betaln(x[..., 0], x[..., 1])

def _max_mask_non_finite(x, axis=-1, keepdims=False, mask=0):
  """Returns `max` or `mask` if `max` is not finite."""
  m = np.max(x, axis=_astuple(axis), keepdims=keepdims)
  needs_masking = ~np.isfinite(m)
  if needs_masking.ndim > 0:
    m[needs_masking] = mask
  elif needs_masking:
    m = mask
  return m

def _softmax(logits, axis=None, name=None):  # pylint: disable=unused-argument
  axis = -1 if axis is None else axis
  y = logits - _max_mask_non_finite(logits, axis=axis, keepdims=True)
  np.exp(y, out=y)
  y /= np.sum(y, axis=_astuple(axis), keepdims=True)
  return y

def _reduce_logsumexp(input_tensor, axis=None, keepdims=False, name=None):  # pylint: disable=unused-argument
  """Computes `log(sum(exp(input_tensor))) along the specified axis."""
  try:
    return scipy_special.logsumexp(
        input_tensor, axis=_astuple(axis), keepdims=keepdims)
  except NotImplementedError:
    # We offer a non SP version just in case SP isn't installed and this
    # because logsumexp is often used.
    m = _max_mask_non_finite(input_tensor, axis=axis, keepdims=True)
    y = input_tensor - m
    y = np.exp(y, out=y)
    return m + np.log(np.sum(y, axis=_astuple(axis), keepdims=keepdims))

# --- Begin Public Functions --------------------------------------------------

abs = utils.copy_docstring(  # pylint: disable=redefined-builtin
    tf.math.abs,
    lambda x, name=None: np.abs(x))

accumulate_n = utils.copy_docstring(
    tf.math.accumulate_n,
    lambda inputs, shape=None, tensor_dtype=None, name=None: (  # pylint: disable=g-long-lambda
        sum(map(np.array, inputs)).astype(utils.numpy_dtype(tensor_dtype))))

acos = utils.copy_docstring(
    tf.math.acos,
    lambda x, name=None: np.arccos(x))

acosh = utils.copy_docstring(
    tf.math.acosh,
    lambda x, name=None: np.arccosh(x))

add = utils.copy_docstring(
    tf.math.add,
    lambda x, y, name=None: np.add(x, y))

add_n = utils.copy_docstring(
    tf.math.add_n,
    lambda inputs, name=None: sum(map(np.array, inputs)))

angle = utils.copy_docstring(
    tf.math.angle,
    lambda input, name=None: np.angle(input))

argmax = utils.copy_docstring(
    tf.math.argmax,
    lambda input, axis=None, output_type=tf.int64, name=None: (  # pylint: disable=g-long-lambda
        np.argmax(input, axis=0 if axis is None else _astuple(axis))
        .astype(utils.numpy_dtype(output_type))))

argmin = utils.copy_docstring(
    tf.math.argmin,
    lambda input, axis=None, output_type=tf.int64, name=None: (  # pylint: disable=g-long-lambda
        np.argmin(input, axis=0 if axis is None else _astuple(axis))
        .astype(utils.numpy_dtype(output_type))))

asin = utils.copy_docstring(
    tf.math.asin,
    lambda x, name=None: np.arcsin(x))

asinh = utils.copy_docstring(
    tf.math.asinh,
    lambda x, name=None: np.arcsinh(x))

atan = utils.copy_docstring(
    tf.math.atan,
    lambda x, name=None: np.arctan(x))

atan2 = utils.copy_docstring(
    tf.math.atan2,
    lambda y, x, name=None: np.arctan2(y, x))

atanh = utils.copy_docstring(
    tf.math.atanh,
    lambda x, name=None: np.arctanh(x))

bessel_i0 = utils.copy_docstring(
    tf.math.bessel_i0,
    lambda x, name=None: scipy_special.i0(x))

bessel_i0e = utils.copy_docstring(
    tf.math.bessel_i0e,
    lambda x, name=None: scipy_special.i0e(x))

bessel_i1 = utils.copy_docstring(
    tf.math.bessel_i1,
    lambda x, name=None: scipy_special.i1(x))

bessel_i1e = utils.copy_docstring(
    tf.math.bessel_i1e,
    lambda x, name=None: scipy_special.i1e(x))

betainc = utils.copy_docstring(
    tf.math.betainc,
    lambda a, b, x, name=None: scipy_special.betainc(a, b, x))

bincount = utils.copy_docstring(
    tf.math.bincount,
    _bincount)

ceil = utils.copy_docstring(
    tf.math.ceil,
    lambda x, name=None: np.ceil(x))

# confusion_matrix = utils.copy_docstring(
#     tf.math.confusion_matrix,
#     lambda labels, predictions, num_classes=None, weights=None,
#     dtype=tf.int32, name=None: ...)

conj = utils.copy_docstring(
    tf.math.conj,
    lambda x, name=None: np.conj(x))

cos = utils.copy_docstring(
    tf.math.cos,
    lambda x, name=None: np.cos(x))

cosh = utils.copy_docstring(
    tf.math.cosh,
    lambda x, name=None: np.cosh(x))

count_nonzero = utils.copy_docstring(
    tf.math.count_nonzero,
    lambda input, axis=None, keepdims=None, dtype=tf.int64, name=None: (  # pylint: disable=g-long-lambda
        np.cast[utils.numpy_dtype(dtype)](np.count_nonzero(input, axis))))

cumprod = utils.copy_docstring(
    tf.math.cumprod,
    lambda x, axis=0, exclusive=False, reverse=False, name=None: (  # pylint: disable=g-long-lambda
        np.cumprod(x, _astuple(axis))))

cumsum = utils.copy_docstring(
    tf.math.cumsum,
    lambda x, axis=0, exclusive=False, reverse=False, name=None: (  # pylint: disable=g-long-lambda
        np.cumsum(x, _astuple(axis))))

digamma = utils.copy_docstring(
    tf.math.digamma,
    lambda x, name=None: scipy_special.digamma(x))

divide = utils.copy_docstring(
    tf.math.divide,
    lambda x, y, name=None: np.divide(x, y))

divide_no_nan = utils.copy_docstring(
    tf.math.divide_no_nan,
    lambda x, y, name=None: np.where(  # pylint: disable=g-long-lambda
        np.broadcast_to(np.equal(y, 0.), np.array(x).shape),
        np.zeros_like(np.divide(x, y)),
        np.divide(x, y)))

equal = utils.copy_docstring(
    tf.math.equal,
    lambda x, y, name=None: np.equal(x, y))

erf = utils.copy_docstring(
    tf.math.erf,
    lambda x, name=None: scipy_special.erf(x))

erfc = utils.copy_docstring(
    tf.math.erfc,
    lambda x, name=None: scipy_special.erfc(x))

exp = utils.copy_docstring(
    tf.math.exp,
    lambda x, name=None: np.exp(x))

expm1 = utils.copy_docstring(
    tf.math.expm1,
    lambda x, name=None: np.expm1(x))

floor = utils.copy_docstring(
    tf.math.floor,
    lambda x, name=None: np.floor(x))

floordiv = utils.copy_docstring(
    tf.math.floordiv,
    lambda x, y, name=None: np.floor_divide(x, y))

greater = utils.copy_docstring(
    tf.math.greater,
    lambda x, y, name=None: np.greater(x, y))

greater_equal = utils.copy_docstring(
    tf.math.greater_equal,
    lambda x, y, name=None: np.greater_equal(x, y))

igamma = utils.copy_docstring(
    tf.math.igamma,
    lambda a, x, name=None: scipy_special.gammainc(a, x))

igammac = utils.copy_docstring(
    tf.math.igammac,
    lambda a, x, name=None: scipy_special.gammaincc(a, x))

imag = utils.copy_docstring(
    tf.math.imag,
    lambda input, name=None: np.imag(input))

# in_top_k = utils.copy_docstring(
#     tf.math.in_top_k,
#     lambda targets, predictions, k, name=None: np.in_top_k)

# invert_permutation = utils.copy_docstring(
#     tf.math.invert_permutation,
#     lambda x, name=None: np.invert_permutation)

is_finite = utils.copy_docstring(
    tf.math.is_finite,
    lambda x, name=None: np.isfinite(x))

is_inf = utils.copy_docstring(
    tf.math.is_inf,
    lambda x, name=None: np.isinf(x))

is_nan = utils.copy_docstring(
    tf.math.is_nan,
    lambda x, name=None: np.isnan(x))

# is_non_decreasing = utils.copy_docstring(
#    tf.math.is_non_decreasing,
#    lambda x, name=None: np.is_non_decreasing)

# is_strictly_increasing = utils.copy_docstring(
#     tf.math.is_strictly_increasing,
#     lambda x, name=None: np.is_strictly_increasing)

# l2_normalize = utils.copy_docstring(
#     tf.math.l2_normalize,
#     lambda x, axis=None, epsilon=1e-12, name=None: np.l2_normalize)

lbeta = utils.copy_docstring(
    tf.math.lbeta,
    _lbeta)

less = utils.copy_docstring(
    tf.math.less,
    lambda x, y, name=None: np.less(x, y))

less_equal = utils.copy_docstring(
    tf.math.less_equal,
    lambda x, y, name=None: np.less_equal(x, y))

lgamma = utils.copy_docstring(
    tf.math.lgamma,
    lambda x, name=None: real(scipy_special.loggamma(x)))

log = utils.copy_docstring(
    tf.math.log,
    lambda x, name=None: np.log(x))

log1p = utils.copy_docstring(
    tf.math.log1p,
    lambda x, name=None: np.log1p(x))

log_sigmoid = utils.copy_docstring(
    tf.math.log_sigmoid,
    lambda x, name=None: -np.log1p(np.exp(-x)))

log_softmax = utils.copy_docstring(
    tf.math.log_softmax,
    lambda logits, axis=None, name=None: (np.subtract(  # pylint: disable=g-long-lambda
        logits,
        reduce_logsumexp(logits, -1 if axis is None else axis, keepdims=True))))

logical_and = utils.copy_docstring(
    tf.math.logical_and,
    lambda x, y, name=None: np.logical_and(x, y))

logical_not = utils.copy_docstring(
    tf.math.logical_not,
    lambda x, name=None: np.logical_not(x))

logical_or = utils.copy_docstring(
    tf.math.logical_or,
    lambda x, y, name=None: np.logical_or(x, y))

logical_xor = utils.copy_docstring(
    tf.math.logical_xor,
    lambda x, y, name=None: np.logical_xor(x, y))

maximum = utils.copy_docstring(
    tf.math.maximum,
    lambda x, y, name=None: np.maximum(x, y))

minimum = utils.copy_docstring(
    tf.math.minimum,
    lambda x, y, name=None: np.minimum(x, y))

multiply = utils.copy_docstring(
    tf.math.multiply,
    lambda x, y, name=None: np.multiply(x, y))

multiply_no_nan = utils.copy_docstring(
    tf.math.multiply_no_nan,
    lambda x, y, name=None: np.where(  # pylint: disable=g-long-lambda
        np.broadcast_to(np.equal(y, 0.), np.array(x).shape),
        np.zeros_like(np.multiply(x, y)),
        np.multiply(x, y)))

negative = utils.copy_docstring(
    tf.math.negative,
    lambda x, name=None: np.negative(x))

# nextafter = utils.copy_docstring(
#     tf.math.nextafter,
#     lambda x1, x2, name=None: np.nextafter)

not_equal = utils.copy_docstring(
    tf.math.not_equal,
    lambda x, y, name=None: np.not_equal(x, y))

polygamma = utils.copy_docstring(
    tf.math.polygamma,
    lambda a, x, name=None: scipy_special.polygamma(a, x))

polyval = utils.copy_docstring(
    tf.math.polyval,
    lambda coeffs, x, name=None: np.polyval(coeffs, x))

pow = utils.copy_docstring(  # pylint: disable=redefined-builtin
    tf.math.pow,
    lambda x, y, name=None: np.power(x, y))

real = utils.copy_docstring(
    tf.math.real,
    lambda input, name=None: np.real(input))

reciprocal = utils.copy_docstring(
    tf.math.reciprocal,
    lambda x, name=None: np.reciprocal(x))

reduce_all = utils.copy_docstring(
    tf.math.reduce_all,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.all(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_any = utils.copy_docstring(
    tf.math.reduce_any,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.any(input_tensor, _astuple(axis), keepdims=keepdims)))

# reduce_euclidean_norm = utils.copy_docstring(
#     tf.math.reduce_euclidean_norm,
#     lambda input_tensor, axis=None, keepdims=False, name=None: (
#         np.reduce_euclidean_norm))

reduce_logsumexp = utils.copy_docstring(
    tf.math.reduce_logsumexp,
    _reduce_logsumexp)

reduce_max = utils.copy_docstring(
    tf.math.reduce_max,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.max(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_mean = utils.copy_docstring(
    tf.math.reduce_mean,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.mean(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_min = utils.copy_docstring(
    tf.math.reduce_min,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.min(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_prod = utils.copy_docstring(
    tf.math.reduce_prod,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.prod(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_std = utils.copy_docstring(
    tf.math.reduce_std,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.std(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_sum = utils.copy_docstring(
    tf.math.reduce_sum,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.sum(input_tensor, _astuple(axis), keepdims=keepdims)))

reduce_variance = utils.copy_docstring(
    tf.math.reduce_variance,
    lambda input_tensor, axis=None, keepdims=False, name=None: (  # pylint: disable=g-long-lambda
        np.var(input_tensor, _astuple(axis), keepdims=keepdims)))

rint = utils.copy_docstring(
    tf.math.rint,
    lambda x, name=None: np.rint(x))

round = utils.copy_docstring(  # pylint: disable=redefined-builtin
    tf.math.round,
    lambda x, name=None: np.round(x))

rsqrt = utils.copy_docstring(
    tf.math.rsqrt,
    lambda x, name=None: 1. / np.sqrt(x))

# scalar_mul = utils.copy_docstring(
#     tf.math.scalar_mul,
#     lambda data, segment_ids, name=None: np.scalar_mul)

# segment_max = utils.copy_docstring(
#     tf.math.segment_max,
#     lambda data, segment_ids, name=None: np.segment_max)

# segment_mean = utils.copy_docstring(
#     tf.math.segment_mean,
#     lambda data, segment_ids, name=None: np.segment_mean)

# segment_min = utils.copy_docstring(
#     tf.math.segment_min,
#     lambda data, segment_ids, name=None: np.segment_min)

# segment_prod = utils.copy_docstring(
#     tf.math.segment_prod,
#     lambda data, segment_ids, name=None: np.segment_prod)

# segment_sum = utils.copy_docstring(
#     tf.math.segment_sum,
#     lambda data, segment_ids, name=None: np.segment_sum)

sigmoid = utils.copy_docstring(
    tf.math.sigmoid,
    lambda x, name=None: 1. / (1. + np.exp(-x)))

sign = utils.copy_docstring(
    tf.math.sign,
    lambda x, name=None: np.sign(x))

sin = utils.copy_docstring(
    tf.math.sin,
    lambda x, name=None: np.sin(x))

sinh = utils.copy_docstring(
    tf.math.sinh,
    lambda x, name=None: np.sinh(x))

softmax = utils.copy_docstring(
    tf.math.softmax,
    _softmax)

softplus = utils.copy_docstring(
    tf.math.softplus,
    lambda x, name=None: np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0.))

softsign = utils.copy_docstring(
    tf.math.softsign,
    lambda features, name=None: np.divide(features, (np.abs(features) + 1)))

sqrt = utils.copy_docstring(
    tf.math.sqrt,
    lambda x, name=None: np.sqrt(x))

square = utils.copy_docstring(
    tf.math.square,
    lambda x, name=None: np.square(x))

squared_difference = utils.copy_docstring(
    tf.math.squared_difference,
    lambda x, y, name=None: np.square(x - y))

subtract = utils.copy_docstring(
    tf.math.subtract,
    lambda x, y, name=None: np.subtract(x, y))

tan = utils.copy_docstring(
    tf.math.tan,
    lambda x, name=None: np.tan(x))

tanh = utils.copy_docstring(
    tf.math.tanh,
    lambda x, name=None: np.tanh(x))

# top_k = utils.copy_docstring(
#     tf.math.top_k,
#     lambda input, k=1, sorted=True, name=None: np.top_k)

truediv = utils.copy_docstring(
    tf.math.truediv,
    lambda x, y, name=None: np.true_divide(x, y))

# unsorted_segment_max = utils.copy_docstring(
#     tf.math.unsorted_segment_max,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_max))

# unsorted_segment_mean = utils.copy_docstring(
#     tf.math.unsorted_segment_mean,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_mean))

# unsorted_segment_min = utils.copy_docstring(
#     tf.math.unsorted_segment_min,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_min))

# unsorted_segment_prod = utils.copy_docstring(
#     tf.math.unsorted_segment_prod,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_prod))

# unsorted_segment_sqrt_n = utils.copy_docstring(
#     tf.math.unsorted_segment_sqrt_n,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_sqrt_n))

# unsorted_segment_sum = utils.copy_docstring(
#     tf.math.unsorted_segment_sum,
#     lambda data, segment_ids, num_segments, name=None: (
#         np.unsorted_segment_sum))

xdivy = utils.copy_docstring(
    tf.math.xdivy,
    lambda x, y, name=None: (  # pylint: disable=unused-argument,g-long-lambda
        np.where(np.equal(x, 0.),
                 np.zeros_like(np.multiply(x, y)),
                 np.divide(x, y))))

xlogy = utils.copy_docstring(
    tf.math.xlogy,
    lambda x, y, name=None: (  # pylint: disable=unused-argument,g-long-lambda
        np.where(np.equal(x, 0.),
                 np.zeros_like(np.multiply(x, y)),
                 np.multiply(x, np.log(y)))))

# zero_fraction = utils.copy_docstring(
#     tf.math.zero_fraction,
#     lambda value, name=None: np.zero_fraction)

zeta = utils.copy_docstring(
    tf.math.zeta,
    lambda x, q, name=None: scipy_special.zeta(x, q))


# 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.
# ============================================================================
"""TF assertions."""

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

# Dependency imports
import numpy as np

import tensorflow as tf

assert_equal = tf.compat.v1.assert_equal
assert_greater = tf.compat.v1.assert_greater
assert_less = tf.compat.v1.assert_less
assert_rank = tf.compat.v1.assert_rank

assert_greater_equal = tf.compat.v1.assert_greater_equal
assert_integer = tf.compat.v1.assert_integer
assert_less_equal = tf.compat.v1.assert_less_equal
assert_near = tf.compat.v1.assert_near
assert_non_negative = tf.compat.v1.assert_non_negative
assert_non_positive = tf.compat.v1.assert_non_positive
assert_none_equal = tf.compat.v1.assert_none_equal
assert_positive = tf.compat.v1.assert_positive
assert_positive = tf.compat.v1.assert_positive
assert_rank_at_least = tf.compat.v1.assert_rank_at_least

def assert_finite(x, data=None, summarize=None, message=None, name=None):
  """Assert all elements of `x` are finite.

  Args:
    x:  Numeric `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_finite".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_finite'):
    x_ = tf.get_static_value(x)
    if x_ is not None:
      if ~np.all(np.isfinite(x_)):
        raise ValueError(message)
      return x
    assertion = tf.compat.v1.assert_equal(
        tf.math.is_finite(x), tf.ones_like(x, tf.bool),
        data=data, summarize=summarize, message=message)
    with tf.control_dependencies([assertion]):
      return tf.identity(x)

def assert_rank_at_most(x, rank, data=None, summarize=None, message=None,
                        name=None):
  """Assert `x` has rank equal to `rank` or smaller.

  Example of adding a dependency to an operation:

  ```python
  with tf.control_dependencies([tf.assert_rank_at_most(x, 2)]):
    output = tf.reduce_sum(x)
  ```

  Args:
    x:  Numeric `Tensor`.
    rank:  Scalar `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_rank_at_most".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_rank_at_most'):
    return tf.compat.v1.assert_less_equal(
        tf.rank(x), rank, data=data, summarize=summarize, message=message)