code: def _extract_log_probs(num_states, dist, move_time_dimension=False):
      """Tabulate log probabilities from a batch of distributions.
  
      Args:
        num_states: Number of states for the hidden variable.
        dist: distribution to get log probabilities from.
        move_time_dimension: optional boolean.
          if True, moves the time dimension to the first dimension.
          Note that the function assumes that the second rightmost dimension
          in dist.batch_shape is the time dimension.
  
      Returns:
        log_probabilities tf.Tensor with shape:
          dist.batch_shape + [num_states] if move_time_dimension is False.
          [num_steps] + dist.batch_shape + [num_states] otherwise.
      """
  
      states = tf.reshape(
          tf.range(num_states),
          ps.concat(
              [[num_states], ps.ones_like(dist.batch_shape_tensor())],
              axis=0,
          ),
      )
      log_probabilities = distribution_util.move_dimension(
          dist.log_prob(states), 0, -1
      )
  
      if move_time_dimension:
          log_probabilities = distribution_util.move_dimension(
              log_probabilities, -3, 0
          )
      return log_probabilities
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  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 4 for calling a global function with example usage: _extract_log_probs(num_states, dist, move_time_dimension) and returns: log_probabilities