def NewLogProbs(dist_inputs, actions, log_prob_fun):
  """Given distribution and actions calculate log probs."""
  new_log_probs = log_prob_fun(dist_inputs,
                               actions)
  return new_log_probs
Given distribution and actions calculate log probs.
trax/rl/rl_layers.py