def scan_discounted_sum(sequence, decay, initial_value, reverse=False,
                        sequence_lengths=None, back_prop=True,
                        name="scan_discounted_sum"):
  """Evaluates a cumulative discounted sum along dimension 0.

    ```python
    if reverse = False:
      result[1] = sequence[1] + decay[1] * initial_value
      result[k] = sequence[k] + decay[k] * result[k - 1]
    if reverse = True:
      result[last] = sequence[last] + decay[last] * initial_value
      result[k] = sequence[k] + decay[k] * result[k + 1]
    ```

  Respective dimensions T, B and ... have to be the same for all input tensors.
  T: temporal dimension of the sequence; B: batch dimension of the sequence.

    if sequence_lengths is set then x1 and x2 below are equivalent:
    ```python
    x1 = zero_pad_to_length(
      scan_discounted_sum(
          sequence[:length], decays[:length], **kwargs), length=T)
    x2 = scan_discounted_sum(sequence, decays,
                             sequence_lengths=[length], **kwargs)
    ```

  Args:
    sequence: Tensor of shape `[T, B, ...]` containing values to be summed.
    decay: Tensor of shape `[T, B, ...]` containing decays/discounts.
    initial_value: Tensor of shape `[B, ...]` containing initial value.
    reverse: Whether to process the sum in a reverse order.
    sequence_lengths: Tensor of shape `[B]` containing sequence lengths to be
      (reversed and then) summed.
    back_prop: Whether to backpropagate.
    name: Sets the name_scope for this op.

  Returns:
    Cumulative sum with discount. Same shape and type as `sequence`.
  """
  # Note this can be implemented in terms of cumprod and cumsum,
  # approximately as (ignoring boundary issues and initial_value):
  #
  # cumsum(decay_prods * sequence) / decay_prods
  # where decay_prods = reverse_cumprod(decay)
  #
  # One reason this hasn't been done is that multiplying then dividing again by
  # products of decays isn't ideal numerically, in particular if any of the
  # decays are zero it results in NaNs.
  with tf.name_scope(name, values=[sequence, decay, initial_value]):
    if sequence_lengths is not None:
      # Zero out sequence and decay beyond sequence_lengths.
      with tf.control_dependencies(
          [tf.assert_equal(sequence.shape[0], decay.shape[0])]):
        mask = tf.sequence_mask(sequence_lengths, maxlen=sequence.shape[0],
                                dtype=sequence.dtype)
        mask = tf.transpose(mask)

      # Adding trailing dimensions to mask to allow for broadcasting.
      to_seq = mask.shape.dims + [1] * (sequence.shape.ndims - mask.shape.ndims)
      sequence *= tf.reshape(mask, to_seq)
      to_decay = mask.shape.dims + [1] * (decay.shape.ndims - mask.shape.ndims)
      decay *= tf.reshape(mask, to_decay)

    sequences = [sequence, decay]
    if reverse:
      sequences = [_reverse_seq(s, sequence_lengths) for s in sequences]

    summed = tf.scan(lambda a, x: x[0] + x[1] * a,
                     sequences,
                     initializer=tf.convert_to_tensor(initial_value),
                     parallel_iterations=1,
                     back_prop=back_prop)

    if not back_prop:
      summed = tf.stop_gradient(summed)
    if reverse:
      summed = _reverse_seq(summed, sequence_lengths)
    return summed
Evaluates a cumulative discounted sum along dimension 0.

    ```python
    if reverse = False:
      result[1] = sequence[1] + decay[1] * initial_value
      result[k] = sequence[k] + decay[k] * result[k - 1]
    if reverse = True:
      result[last] = sequence[last] + decay[last] * initial_value
      result[k] = sequence[k] + decay[k] * result[k + 1]
    ```

  Respective dimensions T, B and ... have to be the same for all input tensors.
  T: temporal dimension of the sequence; B: batch dimension of the sequence.

    if sequence_lengths is set then x1 and x2 below are equivalent:
    ```python
    x1 = zero_pad_to_length(
      scan_discounted_sum(
          sequence[:length], decays[:length], **kwargs), length=T)
    x2 = scan_discounted_sum(sequence, decays,
                             sequence_lengths=[length], **kwargs)
    ```

  Args:
    sequence: Tensor of shape `[T, B, ...]` containing values to be summed.
    decay: Tensor of shape `[T, B, ...]` containing decays/discounts.
    initial_value: Tensor of shape `[B, ...]` containing initial value.
    reverse: Whether to process the sum in a reverse order.
    sequence_lengths: Tensor of shape `[B]` containing sequence lengths to be
      (reversed and then) summed.
    back_prop: Whether to backpropagate.
    name: Sets the name_scope for this op.

  Returns:
    Cumulative sum with discount. Same shape and type as `sequence`.
  
trfl/sequence_ops.py