code: def make_categorical(batch_shape, num_classes, dtype=tf.int32):
      logits = -50.0 + tf.random.uniform(
          list(batch_shape) + [num_classes],
          minval=-10,
          maxval=10,
          dtype=tf.float32,
          seed=test_util.test_seed(),
      )
      return categorical.Categorical(
          logits, dtype=dtype, validate_args=True
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 1 for calling a global function with example usage: make_categorical(batch_shape, num_classes, dtype) and returns: categorical