def METHOD_NAME(rewards: types.NestedArray,
                                   discounts: types.NestedArray,
                                   additional_discount: float, n_step: int):
  """Stolen from TransitionAdder."""
  # Give the same tree structure to the n-step return accumulator,
  # n-step discount accumulator, and self.discount, so that they can be
  # iterated in parallel using tree.map_structure.
  rewards, discounts, self_discount = tree_utils.broadcast_structures(
      rewards, discounts, additional_discount)
  flat_rewards = tree.flatten(rewards)
  flat_discounts = tree.flatten(discounts)
  flat_self_discount = tree.flatten(self_discount)
  # Copy total_discount as it is otherwise read-only.
  total_discount = [np.copy(a[0]) for a in flat_discounts]
  # Broadcast n_step_return to have the broadcasted shape of
  # reward * discount.
  n_step_return = [
      np.copy(np.broadcast_to(r[0],
                              np.broadcast(r[0], d).shape))
      for r, d in zip(flat_rewards, total_discount)
  ]
  # NOTE: total_discount will have one less self_discount applied to it than
  # the value of self._n_step. This is so that when the learner/update uses
  # an additional discount we don't apply it twice. Inside the following loop
  # we will apply this right before summing up the n_step_return.
  for i in range(1, n_step):
    for nsr, td, r, d, sd in zip(n_step_return, total_discount, flat_rewards,
                                 flat_discounts, flat_self_discount):
      # Equivalent to: `total_discount *= self._discount`.
      td *= sd
      # Equivalent to: `n_step_return += reward[i] * total_discount`.
      nsr += r[i] * td
      # Equivalent to: `total_discount *= discount[i]`.
      td *= d[i]
  n_step_return = tree.unflatten_as(rewards, n_step_return)
  total_discount = tree.unflatten_as(rewards, total_discount)
  return n_step_return, total_discount