def METHOD_NAME(
    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