@markdaoust it just keeps getting interesting:

  
      
      colab.research.google.com
  
  
    

Google Colaboratory 7



  
  
    
    
  
  


What I exactly wanted:
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})
Counter({0: 64, 1: 64, 2: 64, 3: 64, 4: 64, 5: 64, 6: 64, 7: 64, 8: 64, 9: 64})

Crux of the code:
def dataset_for_class(i):
    i = tf.cast(i, tf.uint8)
    return support_ds.filter(lambda image, label: label == i).repeat()

sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()

support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))
stratified_ds = tf.data.Dataset.range(10).interleave(dataset_for_class, cycle_length=10) 
stratified_ds = stratified_ds.batch(640)

Notes:

Dataset is CIFAR10.
I made sure that the images getting batched are different as you would ntoice in the notebook provided above.