code: def action_log_probs(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.
      """
      return -F.nll_loss(
          F.log_softmax(torch.flatten(policy_logits, 0, -2), dim=-1),
          torch.flatten(actions),
          reduction="none",
      ).view_as(actions)
  
imports: import collections
  import torch
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: action_log_probs(policy_logits, actions) and returns: unknown