sharded_ds.py

import tensorflow as tf
from functools import partial

def get_ds(from_i, to_i):
    ds = tf.data.Dataset.range(from_i, to_i)
    ds = ds.flat_map(lambda x: tf.data.Dataset.range(x*3, (x + 1 )* 3))
    return ds

datasets = [get_ds(i, i+1) for i in range(0, 5)]

for idx, ds in enumerate(datasets):
    print(f'idx {idx}')
    for sample in ds:
        print(f'  {sample}')

"""
idx 0
  0
  1
  2
idx 1
  3
  4
  5
idx 2
  6
  7
  8
idx 3
  9
  10
  11
idx 4
  12
  13
  14
"""