code: # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  import numpy as np
  from scipy import stats
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import test_util
  tfd = tfp.distributions
  
  
  @test_util.test_all_tf_execution_regimes
  class _SkellamTest(object):
  
    def _make_skellam(self,
                      rate1,
                      rate2,
                      validate_args=True,
                      force_probs_to_zero_outside_support=False):
      return tfd.Skellam(
          rate1=rate1,
          rate2=rate2,
          validate_args=validate_args,
          force_probs_to_zero_outside_support=force_probs_to_zero_outside_support)
  
    def testSkellamShape(self):
      rate1 = tf.constant([3.0] * 5, dtype=self.dtype)
      rate2 = tf.constant([3.0] * 4, dtype=self.dtype)[..., tf.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
  
      self.assertAllEqual(self.evaluate(skellam.batch_shape_tensor()), (4, 5))
      self.assertEqual(skellam.batch_shape, tf.TensorShape([4, 5]))
      self.assertAllEqual(self.evaluate(skellam.event_shape_tensor()), [])
      self.assertEqual(skellam.event_shape, tf.TensorShape([]))
  
    def testInvalidLam(self):
      invalid_rate = self.dtype([-.01, 1., 2.])
      valid_rate = self.dtype([1., 2., 3.])
      with self.assertRaisesOpError('Argument `rate1` must be non-negative.'):
        skellam = self._make_skellam(rate1=invalid_rate, rate2=valid_rate)
        self.evaluate(skellam.rate1_parameter())
  
      with self.assertRaisesOpError('Argument `rate2` must be non-negative.'):
        skellam = self._make_skellam(rate1=valid_rate, rate2=invalid_rate)
        self.evaluate(skellam.rate2_parameter())
  
    def testZeroRate(self):
      lam = self.dtype(0.)
      skellam = tfd.Skellam(rate1=lam, rate2=lam, validate_args=True)
      self.assertAllClose(lam, self.evaluate(skellam.rate1))
      self.assertAllClose(lam, self.evaluate(skellam.rate2))
      self.assertAllClose(0., skellam.prob(3.))
      self.assertAllClose(1., skellam.prob(0.))
      self.assertAllClose(0., skellam.log_prob(0.))
  
    def testSkellamLogPmfDiscreteMatchesScipy(self):
      batch_size = 12
      rate1 = np.linspace(1, 12, 12).astype(self.dtype)
      rate2 = np.array([[1.2], [2.3]]).astype(self.dtype)
      x = np.array([-3., -1., 0., 2., 4., 3., 7., 4., 8., 9., 6., 7.],
                   dtype=self.dtype)
      skellam = self._make_skellam(
          rate1=rate1, rate2=rate2,
          force_probs_to_zero_outside_support=True, validate_args=False)
      log_pmf = skellam.log_prob(x)
      self.assertEqual(log_pmf.shape, (2, batch_size))
      self.assertAllClose(
          self.evaluate(log_pmf),
          stats.skellam.logpmf(x, rate1, rate2))
  
      pmf = skellam.prob(x)
      self.assertEqual(pmf.shape, (2, batch_size,))
      self.assertAllClose(
          self.evaluate(pmf),
          stats.skellam.pmf(x, rate1, rate2))
  
    @test_util.numpy_disable_gradient_test
    def testSkellamLogPmfGradient(self):
      batch_size = 6
      rate1 = tf.constant([3.] * batch_size, dtype=self.dtype)
      rate2 = tf.constant([2.7] * batch_size, dtype=self.dtype)
      x = np.array([-1., 2., 3., 4., 5., 6.], dtype=self.dtype)
  
      err = self.compute_max_gradient_error(
          lambda lam: self._make_skellam(  # pylint:disable=g-long-lambda
              rate1=lam, rate2=rate2).log_prob(x), [rate1])
      self.assertLess(err, 7e-4)
  
      err = self.compute_max_gradient_error(
          lambda lam: self._make_skellam(  # pylint:disable=g-long-lambda
              rate1=rate1, rate2=lam).log_prob(x), [rate2])
      self.assertLess(err, 7e-4)
  
    @test_util.numpy_disable_gradient_test
    def testSkellamLogPmfGradientAtZeroPmf(self):
      # Check that the derivative wrt parameter at the zero-prob points is zero.
      batch_size = 6
      rate1 = tf.constant(np.linspace(1, 7, 6), dtype=self.dtype)
      rate2 = tf.constant(np.linspace(9.1, 12.1, 6), dtype=self.dtype)
      x = tf.constant([-2.1, -1.3, -0.5, 0.2, 1.5, 10.5], dtype=self.dtype)
  
      def make_skellam_log_prob(apply_to_second_rate=False):
        def skellam_log_prob(lam):
          return self._make_skellam(
              rate1=rate1 if apply_to_second_rate else lam,
              rate2=lam if apply_to_second_rate else rate2,
              force_probs_to_zero_outside_support=True,
              validate_args=False).log_prob(x)
        return skellam_log_prob
      _, dlog_pmf_dlam = self.evaluate(tfp.math.value_and_gradient(
          make_skellam_log_prob(), rate1))
  
      self.assertEqual(dlog_pmf_dlam.shape, (batch_size,))
      self.assertAllClose(dlog_pmf_dlam, np.zeros([batch_size]))
  
      _, dlog_pmf_dlam = self.evaluate(tfp.math.value_and_gradient(
          make_skellam_log_prob(True), rate2))
  
      self.assertEqual(dlog_pmf_dlam.shape, (batch_size,))
      self.assertAllClose(dlog_pmf_dlam, np.zeros([batch_size]))
  
    def testSkellamMean(self):
      rate1 = np.array([1.0, 3.0, 2.5], dtype=self.dtype)
      rate2 = np.array([5.0, 7.13, 2.56, 41.], dtype=self.dtype)[..., np.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      self.assertEqual(skellam.mean().shape, (4, 3))
      self.assertAllClose(
          self.evaluate(skellam.mean()), stats.skellam.mean(rate1, rate2))
      self.assertAllClose(self.evaluate(skellam.mean()), rate1 - rate2)
  
    def testSkellamVariance(self):
      rate1 = np.array([1.0, 3.0, 2.5], dtype=self.dtype)
      rate2 = np.array([5.0, 7.13, 2.56, 41.], dtype=self.dtype)[..., np.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      self.assertEqual(skellam.variance().shape, (4, 3))
      self.assertAllClose(
          self.evaluate(skellam.variance()), stats.skellam.var(rate1, rate2))
      self.assertAllClose(self.evaluate(skellam.variance()), rate1 + rate2)
  
    def testSkellamStd(self):
      rate1 = np.array([1.0, 3.0, 2.5], dtype=self.dtype)
      rate2 = np.array([5.0, 7.13, 2.56, 41.], dtype=self.dtype)[..., np.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      self.assertEqual(skellam.stddev().shape, (4, 3))
      self.assertAllClose(
          self.evaluate(skellam.stddev()), stats.skellam.std(rate1, rate2))
      self.assertAllClose(self.evaluate(skellam.stddev()), np.sqrt(rate1 + rate2))
  
    def testSkellamSample(self):
      rate1 = self.dtype([2., 3., 4.])
      rate2 = self.dtype([7.1, 3.2])[..., np.newaxis]
      n = int(2e5)
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      samples = skellam.sample(n, seed=test_util.test_seed())
      sample_values = self.evaluate(samples)
      self.assertEqual(samples.shape, (n, 2, 3))
      self.assertEqual(sample_values.shape, (n, 2, 3))
      self.assertAllClose(
          sample_values.mean(axis=0), stats.skellam.mean(rate1, rate2), rtol=.03)
      self.assertAllClose(
          sample_values.var(axis=0), stats.skellam.var(rate1, rate2), rtol=.03)
  
    def testAssertValidSample(self):
      rate1 = np.array([1.0, 3.0, 2.5], dtype=self.dtype)
      rate2 = np.array([2.1, 7.0, 42.5], dtype=self.dtype)
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      with self.assertRaisesOpError('has non-integer components'):
        self.evaluate(skellam.prob([-1.2, 3., 4.2]))
  
    def testSkellamSampleMultidimensionalMean(self):
      rate1 = self.dtype([2., 3., 4., 5., 6.])
      rate2 = self.dtype([7.1, 3.2, 10., 9.])[..., np.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      n = int(2e5)
      samples = skellam.sample(n, seed=test_util.test_seed())
      sample_values = self.evaluate(samples)
      self.assertEqual(samples.shape, (n, 4, 5))
      self.assertEqual(sample_values.shape, (n, 4, 5))
      self.assertAllClose(
          sample_values.mean(axis=0),
          stats.skellam.mean(rate1, rate2), rtol=.04, atol=0)
  
    def testSkellamSampleMultidimensionalVariance(self):
      rate1 = self.dtype([2., 3., 4., 5., 6.])
      rate2 = self.dtype([7.1, 3.2, 10., 9.])[..., np.newaxis]
      skellam = self._make_skellam(rate1=rate1, rate2=rate2)
      n = int(1e5)
      samples = skellam.sample(n, seed=test_util.test_seed())
      sample_values = self.evaluate(samples)
      self.assertEqual(samples.shape, (n, 4, 5))
      self.assertEqual(sample_values.shape, (n, 4, 5))
  
      self.assertAllClose(
          sample_values.var(axis=0),
          stats.skellam.var(rate1, rate2), rtol=.03, atol=0)
  
    @test_util.tf_tape_safety_test
    def testGradientThroughRate(self):
      rate1 = tf.Variable(3.)
      rate2 = tf.Variable(4.)
      dist = self._make_skellam(rate1=rate1, rate2=rate2)
      with tf.GradientTape() as tape:
        loss = -dist.log_prob([1., 2., 4.])
      grad = tape.gradient(loss, dist.trainable_variables)
      self.assertLen(grad, 2)
      self.assertAllNotNone(grad)
  
    def testAssertsNonNegativeRate(self):
      rate1 = tf.Variable([-1., 2., -3.])
      rate2 = tf.Variable([1., 2., 3.])
      self.evaluate([rate1.initializer, rate2.initializer])
      with self.assertRaisesOpError('Argument `rate1` must be non-negative.'):
        dist = self._make_skellam(
            rate1=rate1, rate2=rate2, validate_args=True)
        self.evaluate(dist.sample(seed=test_util.test_seed()))
  
      rate1 = tf.Variable([1., 2., 3.])
      rate2 = tf.Variable([-1., 2., -3.])
      self.evaluate([rate1.initializer, rate2.initializer])
  
      with self.assertRaisesOpError('Argument `rate2` must be non-negative.'):
        dist = self._make_skellam(
            rate1=rate1, rate2=rate2, validate_args=True)
        self.evaluate(dist.sample(seed=test_util.test_seed()))
  
    def testAssertsNonNegativeRateAfterMutation(self):
      rate1 = tf.Variable([1., 2., 3.])
      rate2 = tf.Variable([1., 2., 3.])
      self.evaluate([rate1.initializer, rate2.initializer])
      dist = self._make_skellam(
          rate1=rate1, rate2=rate2, validate_args=True)
      self.evaluate(dist.mean())
      with self.assertRaisesOpError('Argument `rate1` must be non-negative.'):
        with tf.control_dependencies([rate1.assign([1., 2., -3.])]):
          self.evaluate(dist.sample(seed=test_util.test_seed()))
  
      rate1 = tf.Variable([1., 2., 3.])
      rate2 = tf.Variable([1., 2., 3.])
      self.evaluate([rate1.initializer, rate2.initializer])
      dist = self._make_skellam(
          rate1=rate1, rate2=rate2, validate_args=True)
      self.evaluate(dist.mean())
  
      with self.assertRaisesOpError('Argument `rate2` must be non-negative.'):
        with tf.control_dependencies([rate2.assign([1., 2., -3.])]):
          self.evaluate(dist.sample(seed=test_util.test_seed()))
  
  
  @test_util.test_all_tf_execution_regimes
  class SkellamTestFloat32(test_util.TestCase, _SkellamTest):
    dtype = np.float32
  
  
  @test_util.test_all_tf_execution_regimes
  class SkellamTestFloat64(test_util.TestCase, _SkellamTest):
    dtype = np.float64
  
  
  @test_util.test_all_tf_execution_regimes
  class SkellamLogRateTest(_SkellamTest):
  
    def _make_skellam(self,
                      rate1,
                      rate2,
                      validate_args=True,
                      force_probs_to_zero_outside_support=False):
      return tfd.Skellam(
          log_rate1=tf.math.log(rate1),
          log_rate2=tf.math.log(rate2),
          validate_args=validate_args,
          force_probs_to_zero_outside_support=force_probs_to_zero_outside_support)
  
    # No need to worry about the non-negativity of `rate` when using the
    # `log_rate` parameterization.
    def testInvalidLam(self):
      pass
  
    def testAssertsNonNegativeRate(self):
      pass
  
    def testAssertsNonNegativeRateAfterMutation(self):
      pass
  
    # The gradient is not tracked through tf.math.log(rate) in _make_skellam(),
    # so log_rate needs to be defined as a Variable and passed directly.
    @test_util.tf_tape_safety_test
    def testGradientThroughRate(self):
      log_rate1 = tf.Variable(3.)
      log_rate2 = tf.Variable(4.)
      dist = tfd.Skellam(
          log_rate1=log_rate1, log_rate2=log_rate2, validate_args=True)
      with tf.GradientTape() as tape:
        loss = -dist.log_prob([1., 2., 4.])
      grad = tape.gradient(loss, dist.trainable_variables)
      self.assertLen(grad, 2)
      self.assertAllNotNone(grad)
  
  
  if __name__ == '__main__':
    tf.test.main()
   # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Tests for LogNormal."""
  
  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
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.internal import test_util
  
  
  @test_util.test_all_tf_execution_regimes
  class LogNormalTest(test_util.TestCase):
  
    def setUp(self):
      self._rng = np.random.RandomState(123)
  
    def testLogNormalStats(self):
  
      loc = np.float32([3., 1.5])
      scale = np.float32([0.4, 1.1])
      dist = tfd.LogNormal(loc=loc, scale=scale, validate_args=True)
  
      self.assertAllClose(self.evaluate(dist.mean()),
                          np.exp(loc + scale**2 / 2))
      self.assertAllClose(self.evaluate(dist.variance()),
                          (np.exp(scale**2) - 1) * np.exp(2 * loc + scale**2))
      self.assertAllClose(self.evaluate(dist.stddev()),
                          np.sqrt(self.evaluate(dist.variance())))
      self.assertAllClose(self.evaluate(dist.mode()),
                          np.exp(loc - scale**2))
      self.assertAllClose(self.evaluate(dist.entropy()),
                          np.log(scale * np.exp(loc + 0.5) * np.sqrt(2 * np.pi)))
  
    def testLogNormalSample(self):
      loc, scale = 1.5, 0.4
      dist = tfd.LogNormal(loc=loc, scale=scale, validate_args=True)
      samples = self.evaluate(dist.sample(6000, seed=test_util.test_seed()))
      self.assertAllClose(np.mean(samples),
                          self.evaluate(dist.mean()),
                          atol=0.1)
      self.assertAllClose(np.std(samples),
                          self.evaluate(dist.stddev()),
                          atol=0.1)
  
    def testLogNormalPDF(self):
      loc, scale = 1.5, 0.4
      dist = tfd.LogNormal(loc=loc, scale=scale, validate_args=True)
  
      x = np.array([1e-4, 1.0, 2.0], dtype=np.float32)
  
      log_pdf = dist.log_prob(x)
      analytical_log_pdf = -np.log(x * scale * np.sqrt(2 * np.pi)) - (
          np.log(x) - loc)**2 / (2. * scale**2)
  
      self.assertAllClose(self.evaluate(log_pdf), analytical_log_pdf)
  
    def testLogNormalCDF(self):
      loc, scale = 1.5, 0.4
      dist = tfd.LogNormal(loc=loc, scale=scale, validate_args=True)
  
      x = np.array([1e-4, 1.0, 2.0], dtype=np.float32)
  
      cdf = dist.cdf(x)
      analytical_cdf = .5 + .5 * tf.math.erf(
          (np.log(x) - loc) / (scale * np.sqrt(2)))
      self.assertAllClose(self.evaluate(cdf),
                          self.evaluate(analytical_cdf))
  
    def testLogNormalLogNormalKL(self):
      batch_size = 6
      mu_a = np.array([3.0] * batch_size)
      sigma_a = np.array([1.0, 2.0, 3.0, 1.5, 2.5, 3.5])
      mu_b = np.array([-3.0] * batch_size)
      sigma_b = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
  
      ln_a = tfd.LogNormal(loc=mu_a, scale=sigma_a, validate_args=True)
      ln_b = tfd.LogNormal(loc=mu_b, scale=sigma_b, validate_args=True)
  
      kl = tfd.kl_divergence(ln_a, ln_b)
      kl_val = self.evaluate(kl)
  
      normal_a = tfd.Normal(loc=mu_a, scale=sigma_a, validate_args=True)
      normal_b = tfd.Normal(loc=mu_b, scale=sigma_b, validate_args=True)
      kl_expected_from_normal = tfd.kl_divergence(normal_a, normal_b)
  
      kl_expected_from_formula = ((mu_a - mu_b)**2 / (2 * sigma_b**2) + 0.5 * (
          (sigma_a**2 / sigma_b**2) - 1 - 2 * np.log(sigma_a / sigma_b)))
  
      x = ln_a.sample(int(2e5), seed=test_util.test_seed())
      kl_sample = tf.reduce_mean(ln_a.log_prob(x) - ln_b.log_prob(x), axis=0)
      kl_sample_ = self.evaluate(kl_sample)
  
      self.assertEqual(kl.shape, (batch_size,))
      self.assertAllClose(kl_val, kl_expected_from_normal)
      self.assertAllClose(kl_val, kl_expected_from_formula)
      self.assertAllClose(
          kl_expected_from_formula, kl_sample_, atol=0.0, rtol=1e-2)
  
    # TODO(b/144948687) Avoid `nan` at boundary. Ideally we'd do this test:
    # def testPdfAtBoundary(self):
    #   dist = tfd.LogNormal(loc=5., scale=2.)
    #   pdf = self.evaluate(dist.prob(0.))
    #   log_pdf = self.evaluate(dist.log_prob(0.))
    #   self.assertEqual(pdf, 0.)
    #   self.assertAllNegativeInf(log_pdf)
  
    def testAssertValidSample(self):
      dist = tfd.LogNormal(loc=[-3., 1., 4.], scale=2., validate_args=True)
      with self.assertRaisesOpError('Sample must be non-negative.'):
        self.evaluate(dist.cdf([3., -0.2, 1.]))
  
    def testSupportBijectorOutsideRange(self):
      dist = tfd.LogNormal(loc=1., scale=2., validate_args=True)
      with self.assertRaisesOpError('must be greater than or equal to 0'):
        dist.experimental_default_event_space_bijector().inverse(
            [-4.2, -1e-6, -1.3])
  
  if __name__ == '__main__':
    tf.test.main()
   # 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.
  # ============================================================================
  """Tools for processing Tensors."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  
  __all__ = [
      'convert_nonref_to_tensor',
      'discover_trainable_variables',
      'discover_variables',
      'is_module',
      'is_ref',
      'is_trainable_variable',
      'is_variable',
  ]
  
  
  def convert_nonref_to_tensor(value, dtype=None, dtype_hint=None,
                               as_shape_tensor=False, name=None):
    """Converts the given `value` to a `Tensor` if input is nonreference type.
  
    This function converts Python objects of various types to `Tensor` objects
    only if the input has nonreference semantics. Reference semantics are
    characterized by `tensor_util.is_ref` and is any object which is a
    `tf.Variable` or instance of `tf.Module`. This function accepts any input
    which `tf.convert_to_tensor` would also.
  
    Note: This function diverges from default Numpy behavior for `float` and
      `string` types when `None` is present in a Python list or scalar. Rather
      than silently converting `None` values, an error will be thrown.
  
    Args:
      value: An object whose type has a registered `Tensor` conversion function.
      dtype: Optional element type for the returned tensor. If missing, the
        type is inferred from the type of `value`.
      dtype_hint: Optional element type for the returned tensor,
        used when dtype is None. In some cases, a caller may not have a
        dtype in mind when converting to a tensor, so dtype_hint
        can be used as a soft preference.  If the conversion to
        `dtype_hint` is not possible, this argument has no effect.
      as_shape_tensor: Optional boolean when if `True` uses
        `prefer_static.convert_to_shape_tensor` instead of `tf.convert_to_tensor`
        for JAX compatibility.
      name: Optional name to use if a new `Tensor` is created.
  
    Returns:
      tensor: A `Tensor` based on `value`.
  
    Raises:
      TypeError: If no conversion function is registered for `value` to `dtype`.
      RuntimeError: If a registered conversion function returns an invalid value.
      ValueError: If the `value` is a tensor not of given `dtype` in graph mode.
  
  
    #### Examples:
  
    ```python
    from tensorflow_probability.python.internal import tensor_util
  
    x = tf.Variable(0.)
    y = tensor_util.convert_nonref_to_tensor(x)
    x is y
    # ==> True
  
    x = tf.constant(0.)
    y = tensor_util.convert_nonref_to_tensor(x)
    x is y
    # ==> True
  
    x = np.array(0.)
    y = tensor_util.convert_nonref_to_tensor(x)
    x is y
    # ==> False
    tf.is_tensor(y)
    # ==> True
  
    x = tfp.util.DeferredTensor(13.37, lambda x: x)
    y = tensor_util.convert_nonref_to_tensor(x)
    x is y
    # ==> True
    tf.is_tensor(y)
    # ==> True
    tf.equal(y, 13.37)
    # ==> True
    ```
  
    """
    # We explicitly do not use a tf.name_scope to avoid graph clutter.
    if value is None:
      return None
    if is_ref(value):
      if dtype is None:
        return value
      dtype_base = dtype_util.base_dtype(dtype)
      value_dtype_base = dtype_util.base_dtype(value.dtype)
      if dtype_base != value_dtype_base:
        raise TypeError('Mutable type must be of dtype "{}" but is "{}".'.format(
            dtype_util.name(dtype_base), dtype_util.name(value_dtype_base)))
      return value
    if as_shape_tensor:
      return prefer_static.convert_to_shape_tensor(
          value, dtype=dtype, dtype_hint=dtype_hint, name=name)
    return tf.convert_to_tensor(
        value, dtype=dtype, dtype_hint=dtype_hint, name=name)
  
  
  def is_ref(x):
    """Evaluates if the object has reference semantics.
  
    An object is deemed "reference" if it is a `tf.Variable` instance or is
    derived from a `tf.Module` with `dtype` and `shape` properties.
  
    Args:
      x: Any object.
  
    Returns:
      is_ref: Python `bool` indicating input is has nonreference semantics, i.e.,
        is a `tf.Variable` or a `tf.Module` with `dtype` and `shape` properties.
    """
    # TODO(b/134430874): Consider making this recurse through nests, e.g.,
    # `tensor_util.is_ref([tf.Variable(0.), np.array(1.)])`
    # returns True. Note: we'd need to actually create a tf.Module on user's
    # behalf and it would need a `dtype` and `shape`. (I.e., there would be some
    # work to support this.)
    return (
        is_variable(x) or
        (is_module(x) and hasattr(x, 'dtype') and hasattr(x, 'shape'))
    )
  
  
  def is_variable(x):
    """Returns `True` when input is a `tf.Variable`, otherwise `False`."""
    return isinstance(x, tf.Variable)
  
  
  def is_trainable_variable(x):
    """Returns `True` when input is trainable `tf.Variable`, otherwise `False`."""
    return is_variable(x) and getattr(x, 'trainable', False)
  
  
  def is_module(x):
    """Returns `True` when input is a `tf.Module`, otherwise `False`."""
    return isinstance(x, tf.Module)
  
  
  class _Track(tf.Module):
    """Bridge to create functional interface for variable tracking."""
  
    def __init__(self, *args, **kwargs):
      self._args = args
      self._kwargs = kwargs
  
  
  def discover_trainable_variables(x):
    """Returns `tuple` of all trainable `tf.Variables` discoverable in input.
  
    Warning: unlike possibly `tf.Module`, use of this function only does a static,
    "one-time" discovery. (This is self-evidently true from its functional
    nature.)
  
    Args:
      x: An object to inspected for `tf.Variable` dependencies.
  
    Returns:
      trainable_vars: A Python `tuple` of `tf.Variable`s with `trainable=True`.
    """
    return _Track(x).trainable_variables
  
  
  def discover_variables(x):
    """Returns `tuple` of all `tf.Variables` discoverable in input.
  
    Warning: unlike possibly `tf.Module`, use of this function only does a static,
    "one-time" discovery. (This is self-evidently true from its functional
    nature.)
  
    Args:
      x: An object to inspected for `tf.Variable` dependencies.
  
    Returns:
      vars: A Python `tuple` of `tf.Variable`s, regardless of their value of
        `trainable`.
    """
    return _Track(x).variables
  
apis: numpy.sqrt scipy.stats.skellam.std numpy.zeros tensorflow.compat.v2.test.main tensorflow.compat.v2.GradientTape scipy.stats.skellam.mean scipy.stats.skellam.var scipy.stats.skellam.logpmf tensorflow.compat.v2.TensorShape scipy.stats.skellam.pmf numpy.array tensorflow.compat.v2.constant numpy.linspace tensorflow.compat.v2.Variable tensorflow.compat.v2.math.log
apis: numpy.sqrt numpy.float32 numpy.exp tensorflow.compat.v2.test.main numpy.random.RandomState numpy.log numpy.array numpy.std numpy.mean
apis: tensorflow.compat.v2.convert_to_tensor