def log_probs_from_logits_and_actions(policy_logits, actions):
  """Computes action log-probs from policy logits and actions.

  In the notation used throughout documentation and comments, T refers to the
  time dimension ranging from 0 to T-1. B refers to the batch size and
  NUM_ACTIONS refers to the number of actions.

  Args:
    policy_logits: A float32 tensor of shape `[T, B, NUM_ACTIONS]` with
      un-normalized log-probabilities parameterizing a softmax policy.
    actions: An int32 tensor of shape `[T, B]` with actions.

  Returns:
    A float32 tensor of shape `[T, B]` corresponding to the sampling log
    probability of the chosen action w.r.t. the policy.
  """
  policy_logits = tf.convert_to_tensor(policy_logits, dtype=tf.float32)
  actions = tf.convert_to_tensor(actions, dtype=tf.int32)

  policy_logits.shape.assert_has_rank(3)
  actions.shape.assert_has_rank(2)

  return -tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=policy_logits, labels=actions)
Computes action log-probs from policy logits and actions.

  In the notation used throughout documentation and comments, T refers to the
  time dimension ranging from 0 to T-1. B refers to the batch size and
  NUM_ACTIONS refers to the number of actions.

  Args:
    policy_logits: A float32 tensor of shape `[T, B, NUM_ACTIONS]` with
      un-normalized log-probabilities parameterizing a softmax policy.
    actions: An int32 tensor of shape `[T, B]` with actions.

  Returns:
    A float32 tensor of shape `[T, B]` corresponding to the sampling log
    probability of the chosen action w.r.t. the policy.
  
trfl/vtrace_ops.py