code: def testDistribution(self, dist_name, data):
      seed = test_util.test_seed()
      # Explicitly draw event_dim here to avoid relying on _params_event_ndims
      # later, so this test can support distributions that do not implement the
      # slicing protocol.
      event_dim = data.draw(hps.integers(min_value=2, max_value=6))
      dist = data.draw(
          dhps.distributions(
              dist_name=dist_name, event_dim=event_dim, enable_vars=True
          )
      )
      batch_shape = dist.batch_shape
      batch_shape2 = data.draw(
          tfp_hps.broadcast_compatible_shape(batch_shape)
      )
      dist2 = data.draw(
          dhps.distributions(
              dist_name=dist_name,
              batch_shape=batch_shape2,
              event_dim=event_dim,
              enable_vars=True,
          )
      )
      self.evaluate([var.initializer for var in dist.variables])
  
      # Check that the distribution passes Variables through to the accessor
      # properties (without converting them to Tensor or anything like that).
      for k, v in six.iteritems(dist.parameters):
          if not tensor_util.is_ref(v):
              continue
          self.assertIs(getattr(dist, k), v)
  
      if dist_name in SKIP_EXCESSIVE_VAR_USAGE:
          return
  
      # Check that standard statistics do not read distribution parameters more
      # than twice (once in the stat itself and up to once in any validation
      # assertions).
      max_permissible = 2 + extra_tensor_conversions_allowed(dist)
      for stat in sorted(
          data.draw(
              hps.sets(
                  hps.one_of(
                      map(
                          hps.just,
                          [
                              "covariance",
                              "entropy",
                              "mean",
                              "mode",
                              "stddev",
                              "variance",
                          ],
                      )
                  ),
                  min_size=3,
                  max_size=3,
              )
          )
      ):
          hp.note(
              "Testing excessive var usage in {}.{}".format(
                  dist_name, stat
              )
          )
          try:
              with tfp_hps.assert_no_excessive_var_usage(
                  "statistic `{}` of `{}`".format(stat, dist),
                  max_permissible=max_permissible,
              ):
                  getattr(dist, stat)()
  
          except NotImplementedError:
              pass
  
      # Check that `sample` doesn't read distribution parameters more than twice,
      # and that it produces non-None gradients (if the distribution is fully
      # reparameterized).
      with tf.GradientTape() as tape:
          # TDs do bijector assertions twice (once by distribution.sample, and once
          # by bijector.forward).
          max_permissible = 2 + extra_tensor_conversions_allowed(dist)
          with tfp_hps.assert_no_excessive_var_usage(
              "method `sample` of `{}`".format(dist),
              max_permissible=max_permissible,
          ):
              sample = dist.sample(seed=seed)
      if dist.reparameterization_type == tfd.FULLY_REPARAMETERIZED:
          grads = tape.gradient(sample, dist.variables)
          for grad, var in zip(grads, dist.variables):
              var_name = var.name.rstrip("_0123456789:")
              if var_name in NO_SAMPLE_PARAM_GRADS.get(dist_name, ()):
                  continue
              if grad is None:
                  raise AssertionError(
                      "Missing sample -> {} grad for distribution {}".format(
                          var_name, dist_name
                      )
                  )
  
      # Turn off validations, since TODO(b/129271256) log_prob can choke on dist's
      # own samples.  Also, to relax conversion counts for KL (might do >2 w/
      # validate_args).
      dist = dist.copy(validate_args=False)
      dist2 = dist2.copy(validate_args=False)
  
      # Test that KL divergence reads distribution parameters at most once, and
      # that it produces non-None gradients.
      try:
          for d1, d2 in (dist, dist2), (dist2, dist):
              if dist_name in SKIP_KL_CHECK_DIST_VAR_GRADS:
                  continue
              with tf.GradientTape() as tape:
                  with tfp_hps.assert_no_excessive_var_usage(
                      "`kl_divergence` of (`{}` (vars {}), `{}` (vars {}))".format(
                          d1, d1.variables, d2, d2.variables
                      ),
                      max_permissible=1,
                  ):  # No validation => 1 convert per var.
                      kl = d1.kl_divergence(d2)
              wrt_vars = list(d1.variables) + list(d2.variables)
              grads = tape.gradient(kl, wrt_vars)
              for grad, var in zip(grads, wrt_vars):
                  if (
                      grad is None
                      and dist_name not in NO_KL_PARAM_GRADS
                  ):
                      raise AssertionError(
                          "Missing KL({} || {}) -> {} grad:\n"  # pylint: disable=duplicate-string-formatting-argument
                          "{} vars: {}\n{} vars: {}".format(
                              d1,
                              d2,
                              var,
                              d1,
                              d1.variables,
                              d2,
                              d2.variables,
                          )
                      )
      except NotImplementedError:
          # Raised by kl_divergence if no registered KL is found.
          pass
  
      # Test that log_prob produces non-None gradients, except for distributions
      # on the NO_LOG_PROB_PARAM_GRADS blocklist.
      if dist_name not in NO_LOG_PROB_PARAM_GRADS:
          with tf.GradientTape() as tape:
              lp = dist.log_prob(tf.stop_gradient(sample))
          grads = tape.gradient(lp, dist.variables)
          for grad, var in zip(grads, dist.variables):
              if grad is None:
                  raise AssertionError(
                      "Missing log_prob -> {} grad for distribution {}".format(
                          var, dist_name
                      )
                  )
  
      # Test that all forms of probability evaluation avoid reading distribution
      # parameters more than once.
      for evaluative in sorted(
          data.draw(
              hps.sets(
                  hps.one_of(
                      map(
                          hps.just,
                          [
                              "log_prob",
                              "prob",
                              "log_cdf",
                              "cdf",
                              "log_survival_function",
                              "survival_function",
                          ],
                      )
                  ),
                  min_size=3,
                  max_size=3,
              )
          )
      ):
          hp.note(
              "Testing excessive var usage in {}.{}".format(
                  dist_name, evaluative
              )
          )
          try:
              # No validation => 1 convert. But for TD we allow 2:
              # dist.log_prob(bijector.inverse(samp)) + bijector.ildj(samp)
              max_permissible = 2 + extra_tensor_conversions_allowed(
                  dist
              )
              with tfp_hps.assert_no_excessive_var_usage(
                  "evaluative `{}` of `{}`".format(evaluative, dist),
                  max_permissible=max_permissible,
              ):
                  getattr(dist, evaluative)(sample)
          except NotImplementedError:
              pass
  
imports: import collections
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python import distributions
  from tensorflow_probability.python.distributions import hypothesis_testlib
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the DistributionGradientTapeAndConcretizationTest obj's underlying member method code with example usage: obj.testDistribution(dist_name, data) without return types