code: def testCompareToCategorical(self):
      # Use the same probabilities for normal categorical and shifted one.
      shift = 2
      probabilities = [0.3, 0.3, 0.4]
      distribution = tfp.distributions.Categorical(probs=probabilities)
      shifted_distribution = shifted_categorical.ShiftedCategorical(
          probs=probabilities, shift=shift
      )
  
      # Compare outputs of basic methods, using the same starting seed.
      tf.compat.v1.set_random_seed(
          1
      )  # required per b/131171329, only with TF2.
      sample = distribution.sample(seed=1)
      tf.compat.v1.set_random_seed(
          1
      )  # required per b/131171329, only with TF2.
      shifted_sample = shifted_distribution.sample(seed=1)
  
      mode = distribution.mode()
      shifted_mode = shifted_distribution.mode()
  
      sample, shifted_sample = self.evaluate([sample, shifted_sample])
      mode, shifted_mode = self.evaluate([mode, shifted_mode])
  
      self.assertEqual(shifted_sample, sample + shift)
      self.assertEqual(shifted_mode, mode + shift)
  
      # These functions should return the same values for shifted values.
      fns = ["cdf", "log_cdf", "prob", "log_prob"]
      for fn_name in fns:
          fn = getattr(distribution, fn_name)
          shifted_fn = getattr(shifted_distribution, fn_name)
          value, shifted_value = self.evaluate(
              [fn(sample), shifted_fn(shifted_sample)]
          )
          self.assertEqual(value, shifted_value)
  
file_path: agents/tf_agents/distributions/shifted_categorical_test.py
imports: import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import shifted_categorical
label_desc: Use image node_id 3 for calling the ShiftedCategoricalTest obj's underlying member method code with example usage: obj.testCompareToCategorical() without return types