simple_window.py

import tensorflow as tf
from functools import partial

def collate_pair(x, window):
    # a, last sample of b
    return x[:, 0], [x[-1, 1]]

ds_raw = tf.constant([[i, 1 if i % 3 == 0 else 0] for i in range(0, 100)])
ds = tf.data.Dataset.from_tensor_slices(ds_raw)

ds = ds.window(2).flat_map(lambda x: x.batch(2)).map(partial(collate_pair, window=2))

for a, b in ds.take(1):
    print(f'a: {a} -> b: {b}')
    
# a: [0 1] -> b: [0]