code: def _compute_n_step_sequence_targets(
      r_t: tf.Tensor,
      pcont_t: tf.Tensor,
      bootstrap_value: tf.Tensor,
      n: int,
  ) -> tf.Tensor:
      """Computes n-step bootstrapped returns over a sequence.
  
      Args:
        r_t: 2-D tensor of shape [T, B] corresponding to rewards.
        pcont_t: 2-D tensor of shape [T, B] corresponding to pcontinues.
        bootstrap_value: 2-D tensor of shape [T, B] corresponding to bootstrap
          values.
        n: number of steps over which to accumulate reward before bootstrapping.
  
      Returns:
        2-D tensor of shape [T, B] corresponding to bootstrapped returns.
      """
      time_size, batch_size = r_t.shape.as_list()
  
      # Pad r_t and pcont_t so we can use static slice shapes in scan.
      r_t = tf.concat([r_t, tf.zeros((n - 1, batch_size))], 0)
      pcont_t = tf.concat([pcont_t, tf.ones((n - 1, batch_size))], 0)
  
      # We need to use tf.slice with static shapes for TPU compatibility.
      def _slice(tensor, index, size):
          return tf.slice(tensor, [index, 0], [size, batch_size])
  
      # Construct correct bootstrap targets for each time slice t, which are exactly
      # the target values at timestep min(t+n-1, time_size-1).
      last_bootstrap_value = _slice(bootstrap_value, time_size - 1, 1)
      if time_size > n - 1:
          full_bootstrap_steps = [
              _slice(bootstrap_value, n - 1, time_size - (n - 1))
          ]
          truncated_bootstrap_steps = [last_bootstrap_value] * (n - 1)
      else:
          # Only truncated steps, since n > time_size.
          full_bootstrap_steps = []
          truncated_bootstrap_steps = [last_bootstrap_value] * time_size
      bootstrap_value = tf.concat(
          full_bootstrap_steps + truncated_bootstrap_steps, 0
      )
  
      # Iterate backwards for n steps to construct n-step return targets.
      targets = bootstrap_value
      for i in range(n - 1, -1, -1):
          this_pcont_t = _slice(pcont_t, i, time_size)
          this_r_t = _slice(r_t, i, time_size)
          targets = this_r_t + this_pcont_t * targets
      return targets
  
file_path: acme/acme/tf/losses/r2d2.py
imports: from typing import Iterable, NamedTuple, Sequence
  import tensorflow
  import trfl
label_desc: Use image node_id 2 for calling a global function with example usage: _compute_n_step_sequence_targets(r_t, pcont_t, bootstrap_value, n) and returns: targets, tf