code: def discounted_returns(rewards, gammas):
      """Computes discounted returns for a trajectory or a batch of them."""
      returns = np.zeros_like(rewards)
      ret = 0.0
      for i in reversed(range(rewards.shape[-1])):
          ret = rewards[..., i] + gammas[..., i] * ret
          returns[..., i] = ret
      return returns
  
file_path: trax/trax/rl/advantages.py
imports: import gin
  import numpy
  from trax import fastmath
label_desc: Use image node_id 2 for calling a global function with example usage: discounted_returns(rewards, gammas) and returns: returns