code: 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
  
file_path: trax/trax/rl/rl_layers.py
imports: from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: NewLogProbs(dist_inputs, actions, log_prob_fun) and returns: new_log_probs