tf_rolling_window.py

def rolling_window(tensor, dtype, shape, capacity=None):
  with tf.name_scope('rolling_window'):
    window_size = shape[0]
    if capacity is None:
      capacity = shape[0] * 2 + 1 

    q = tf.FIFOQueue(capacity, [dtype], shapes=[shape[1:]])
    enqueue = q.enqueue_many(tensor)
    tf.train.add_queue_runner(
        tf.train.QueueRunner(queue=q, enqueue_ops=[enqueue])
    ) 

    # Pad first element as it will be immediately overwritten
    window_initial_value = q.dequeue_many(window_size - 1)
    window_initial_value_padded = tf.concat(0, [ 
        [window_initial_value[0]],
        window_initial_value
    ])

    window = tf.Variable(
        window_initial_value_padded,
        trainable=False
    ) 
    oldest_pos = tf.Variable(-1, trainable=False)

    updated_pos = oldest_pos.assign((oldest_pos + 1) % window_size)
    updated_window = window[updated_pos].assign(q.dequeue())

    return tf.concat(0, [ 
      updated_window[updated_pos+1:],
      updated_window[:updated_pos+1]
    ])