class_docstr: OneHotCategorical distribution.
  
  The categorical distribution is parameterized by the log-probabilities
  of a set of classes. The difference between OneHotCategorical and Categorical
  distributions is that OneHotCategorical is a discrete distribution over
  one-hot bit vectors whereas Categorical is a discrete distribution over
  positive integers. OneHotCategorical is equivalent to Categorical except
  Categorical has event_dim=() while OneHotCategorical has event_dim=K, where
  K is the number of classes.
  
  This class provides methods to create indexed batches of OneHotCategorical
  distributions. If the provided `logits` or `probs` is rank 2 or higher, for
  every fixed set of leading dimensions, the last dimension represents one
  single OneHotCategorical distribution. When calling distribution
  functions (e.g. `dist.prob(x)`), `logits` and `x` are broadcast to the
  same shape (if possible). In all cases, the last dimension of `logits,x`
  represents single OneHotCategorical distributions.
  
  #### Examples
  
  Creates a 3-class distribution, with the 2nd class, the most likely to be
  drawn from.
  
  ```python
  p = [0.1, 0.5, 0.4]
  dist = OneHotCategorical(probs=p)
  ```
  
  Creates a 3-class distribution, with the 2nd class the most likely to be
  drawn from, using logits.
  
  ```python
  logits = [-2, 2, 0]
  dist = OneHotCategorical(logits=logits)
  ```
  
  Creates a 3-class distribution, with the 3rd class is most likely to be drawn.
  
  ```python
  # counts is a scalar.
  p = [0.1, 0.4, 0.5]
  dist = OneHotCategorical(probs=p)
  dist.prob([0,1,0])  # Shape []
  
  # p will be broadcast to [[0.1, 0.4, 0.5], [0.1, 0.4, 0.5]] to match.
  samples = [[0,1,0], [1,0,0]]
  dist.prob(samples)  # Shape [2]
  ```
code: def _log_prob(self, x):
      logits = self._logits_parameter_no_checks()
      event_size = self._event_size(logits)
  
      x = tf.cast(x, logits.dtype)
  
      # broadcast logits or x if need be.
      if (
          not tensorshape_util.is_fully_defined(x.shape)
          or not tensorshape_util.is_fully_defined(logits.shape)
          or x.shape != logits.shape
      ):
          broadcast_shape = ps.broadcast_shape(
              ps.shape(logits), ps.shape(x)
          )
          logits = tf.broadcast_to(logits, broadcast_shape)
          x = tf.broadcast_to(x, broadcast_shape)
  
      logits_shape = ps.shape(tf.reduce_sum(logits, axis=-1))
      logits_2d = tf.reshape(logits, [-1, event_size])
      x_2d = tf.reshape(x, [-1, event_size])
      ret = -tf.nn.softmax_cross_entropy_with_logits(
          labels=tf.stop_gradient(x_2d), logits=logits_2d
      )
  
      # Reshape back to user-supplied batch and sample dims prior to 2D reshape.
      ret = tf.reshape(ret, logits_shape)
      return ret
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 9 for calling the OneHotCategorical obj's underlying member method code with example usage: obj._log_prob(x) and returns: ret