docs: def discounted_return(reward, length, discount):
    """Discounted Monte-Carlo returns."""
    timestep = tf.range(reward.shape[1].value)
    mask = tf.cast(timestep[None, :] < length[:, None], tf.float32)
    return_ = tf.reverse(tf.transpose(tf.scan(
        lambda agg, cur: cur + discount * agg,
        tf.transpose(tf.reverse(mask * reward, [1]), [1, 0]),
        tf.zeros_like(reward[:, -1]), 1, False), [1, 0]), [1])
    return tf.check_numerics(tf.stop_gradient(return_), 'return')