code: def discounted_return(
      rewards,
      discounts,
      final_value=None,
      time_major=True,
      provide_all_returns=True,
  ):
      """Computes discounted return.
  
      ```
      Q_n = sum_{n'=n}^N gamma^(n'-n) * r_{n'} + gamma^(N-n+1)*final_value.
      ```
  
      For details, see
      "Reinforcement Learning: An Introduction" Second Edition
      by Richard S. Sutton and Andrew G. Barto
  
      Define abbreviations:
      `B`: batch size representing number of trajectories.
      `T`: number of steps per trajectory.  This is equal to `N - n` in the equation
           above.
  
      **Note** To replicate the calculation `Q_n` exactly, use
      `discounts = gamma * tf.ones_like(rewards)` and `provide_all_returns=False`.
  
      Args:
        rewards: Tensor with shape `[T, B]` (or `[T]`) representing rewards.
        discounts: Tensor with shape `[T, B]` (or `[T]`) representing discounts.
        final_value: (Optional.).  Default: An all zeros tensor.  Tensor with shape
          `[B]` (or `[1]`) representing value estimate at `T`. This is optional;
          when set, it allows final value to bootstrap the reward computation.
        time_major: A boolean indicating whether input tensors are time major. False
          means input tensors have shape `[B, T]`.
        provide_all_returns: A boolean; if True, this will provide all of the
          returns by time dimension; if False, this will only give the single
          complete discounted return.
  
      Returns:
        If `provide_all_returns`:
          A tensor with shape `[T, B]` (or `[T]`) representing the discounted
          returns. The shape is `[B, T]` when `not time_major`.
        If `not provide_all_returns`:
          A tensor with shape `[B]` (or []) representing the discounted returns.
      """
      if not time_major:
          with tf.name_scope("to_time_major_tensors"):
              discounts = tf.transpose(discounts)
              rewards = tf.transpose(rewards)
  
      if final_value is None:
          final_value = tf.zeros_like(rewards[-1])
  
      def discounted_return_fn(
          accumulated_discounted_reward, reward_discount
      ):
          reward, discount = reward_discount
          return accumulated_discounted_reward * discount + reward
  
      if provide_all_returns:
          returns = tf.nest.map_structure(
              tf.stop_gradient,
              tf.scan(
                  fn=discounted_return_fn,
                  elems=(rewards, discounts),
                  reverse=True,
                  initializer=final_value,
              ),
          )
  
          if not time_major:
              with tf.name_scope("to_batch_major_tensors"):
                  returns = tf.transpose(returns)
      else:
          returns = tf.foldr(
              fn=discounted_return_fn,
              elems=(rewards, discounts),
              initializer=final_value,
              back_prop=False,
          )
  
      return tf.stop_gradient(returns)
  
label_desc: Use image node_id 1 for calling a global function with example usage: discounted_return(rewards, discounts, final_value, time_major, provide_all_returns) and returns: tf, unknown