How to efficiently extract all slices of given length using tensorflow

Question

I am trying to extract all slices of length 4 along 0th axis of a 2-dim tensor. So far I can do it mixing pure Python with tensorflow.

r = test.shape[0] # test should be a tensor
n = 4
a_list = list(range(r))
the_list = np.array([a_list[slice(i, i+n)] for i in range(r - n+1)])
test_stacked = tf.stack(tf.gather(test, the_list))

What would be an efficient way of doing that without using pure Python? Note that the "test" array is actually supposed to be a tensor, thus its shape isn't known before I execute the first part of the graph.

A full vanilla example:

array = np.array([[0, 1],[1, 2],[2, 3],[3, 4],[4, 5],[5, 6]])
array.shape # (6,2)

r = array.shape[0]
n = 4
a_list = list(range(r))
the_list = np.array([a_list[slice(i, i+n)] for i in range(r - n+1)])

result = array[the_list] # all possible slices of length 4 of the array along 0th axis
result.shape # (3, 4, 2)

result:

[[[0 1]
  [1 2]
  [2 3]
  [3 4]]

 [[1 2]
  [2 3]
  [3 4]
  [4 5]]

 [[2 3]
  [3 4]
  [4 5]
  [5 6]]]

Answer 1

You may want to try the more general tf.extract_image_patches.

import tensorflow as tf

a = tf.constant([[0, 1],[1, 2],[2, 3],[3, 4],[4, 5],[5, 6]])
# tf.extract_image_patches requires a [batch, in_rows, in_cols, depth] tensor
a = a[None, :, :, None]
b = tf.extract_image_patches(a,
  ksizes=[1, 4, 2, 1],
  strides=[1, 1, 1, 1],
  rates=[1, 1, 1, 1],
  padding='VALID')
b = tf.reshape(tf.squeeze(b), [-1, 4, 2])

sess = tf.InteractiveSession()
print(b.eval())

Answer 2

I believe gather_nd is what you are looking for.

# a is a tensor of size (6, 2)

def get_indices(l, d):
    return [[[j] for j in range(i, i + d)] for i in range(l - d + 1)]

b = tf.gather_nd(a, get_indices(6, 4))
# b is a tensor of shape (3, 4, 2)