Window Multidimensional Tensorflow Dataset

Question

I have 2-dimensional data with shape m by n that I want to window with size w along the first axis into a dataset of m-w many two-dimensional arrays each of size w by n. For instance if the data is:

[[0,  1,  2 ], 
 [3,  4,  5 ], 
 [6,  7,  8 ],
 [9,  10, 11]]

then I want to window it into

[[[0, 1 , 2 ], 
  [3, 4 , 5 ], 
  [6, 7 , 8 ]],
 [[3, 4 , 5 ],
  [6, 7 , 8 ],
  [9, 10, 11]]]

I can window the data together into the right sets:

dataset = tf.data.Dataset.from_tensor_slices(np.arange(5*3).reshape(5,3))
dataset = dataset.window(size=3,shift=1,drop_remainder=True)
for window in dataset : print(list(window.as_numpy_iterator()))
>>>[array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8])]
>>>[array([3, 4, 5]), array([6, 7, 8]), array([ 9, 10, 11])]
>>>[array([6, 7, 8]), array([ 9, 10, 11]), array([12, 13, 14])]

but I can't figure out how to get the data back into the stacked shape again. I thought maybe tf.stack, but no dice on that. Does anybody know how to finish this?

Answer 1

I found the answer here actually. I don't know why it works, but it does:

dataset = tf.data.Dataset.from_tensor_slices(np.arange(5*3).reshape(5,3))
dataset = dataset.window(size=3,shift=1)
dataset = dataset.flat_map(lambda x : x.batch(3))
for d in dataset : print(d)

which makes

tf.Tensor(
[[0 1 2]
 [3 4 5]
 [6 7 8]], shape=(3, 3), dtype=int64)
tf.Tensor(
[[ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]], shape=(3, 3), dtype=int64)
tf.Tensor(
[[ 6  7  8]
 [ 9 10 11]
 [12 13 14]], shape=(3, 3), dtype=int64)
tf.Tensor(
[[ 9 10 11]
 [12 13 14]], shape=(2, 3), dtype=int64)
tf.Tensor([[12 13 14]], shape=(1, 3), dtype=int64)