code: def testPdf(self):
      temperature = 0.4
      logits = np.array([[0.3, 0.1, 0.4]]).astype(np.float32)
      dist = roc.RelaxedOneHotCategorical(
          temperature, logits, validate_args=True
      )
      x = self.evaluate(dist.sample(seed=test_util.test_seed()))
      pdf = self.evaluate(dist.prob(x))
      expected_pdf = analytical_pdf(x, temperature, logits)
      self.assertAllClose(expected_pdf.flatten(), pdf, rtol=1e-4)
  
      # variable batch size
      logits = np.array([[0.3, 0.1, 0.4], [0.6, -0.1, 2.0]]).astype(
          np.float32
      )
      temperatures = np.array([0.4, 2.3]).astype(np.float32)
      dist = roc.RelaxedOneHotCategorical(
          temperatures, logits, validate_args=True
      )
      x = self.evaluate(dist.sample(seed=test_util.test_seed()))
      pdf = self.evaluate(dist.prob(x))
      expected_pdf = analytical_pdf(x, temperatures, logits)
      self.assertAllClose(expected_pdf.flatten(), pdf, rtol=1e-4)
  
      # broadcast logits over temparatures
      logits = np.array([0.3, 0.1, 0.4]).astype(np.float32)
      temperatures = np.array([0.4, 2.3]).astype(np.float32)
      dist = roc.RelaxedOneHotCategorical(
          temperatures, logits, validate_args=True
      )
      x = self.evaluate(dist.sample(seed=test_util.test_seed()))
      pdf = self.evaluate(dist.prob(x))
      expected_pdf = analytical_pdf(x, temperatures, logits)
      self.assertAllClose(expected_pdf.flatten(), pdf, rtol=1e-4)
  
imports: from absl.testing import parameterized
  import numpy
  from scipy.special import gamma
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import relaxed_onehot_categorical
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 6 for calling the RelaxedOneHotCategoricalTest obj's underlying member method code with example usage: obj.testPdf() without return types