def _extract_log_probs(num_states, dist):
states = tf.reshape(tf.range(num_states),<EOL>tf.concat([[num_states],<EOL>tf.ones_like(dist.batch_shape_tensor())],<EOL>axis=<NUM_LIT:0>))<EOL>return distribution_util.move_dimension(dist.log_prob(states), <NUM_LIT:0>, -<NUM_LIT:1>)<EOL>
Tabulate log probabilities from a batch of distributions.