I tried another approach:
sampled_idx = np.random.choice(len(x_train), 4000)
sampled_train, sampled_labels = x_train[sampled_idx], y_train[sampled_idx].squeeze()
sampled_labels = sampled_labels.astype("int32")
support_ds = tf.data.Dataset.from_tensor_slices((sampled_train, sampled_labels))

ds = []

for i in np.arange(0, 10):
    ds_label =  (
        support_ds
        .filter(lambda image, label: label==i)
        .repeat())
    ds.append(ds_label)

balanced_ds = tf.data.experimental.sample_from_datasets(
    ds, [0.1] * 10).batch(SUPPORT_BATCH_SIZE)

But here also when I do:
for samples, labels in balanced_ds.take(10):
    print(Counter(labels.numpy()))

the distribution does not come out as expected:
Counter({9: 74, 0: 73, 3: 71, 8: 70, 1: 70, 5: 67, 7: 64, 2: 55, 6: 51, 4: 45})
Counter({2: 76, 3: 70, 4: 68, 1: 67, 6: 64, 0: 62, 7: 62, 8: 60, 9: 56, 5: 55})
Counter({1: 78, 2: 75, 7: 74, 0: 68, 9: 67, 3: 61, 5: 58, 8: 55, 4: 54, 6: 50})
Counter({6: 82, 9: 69, 5: 68, 4: 64, 1: 63, 3: 62, 7: 62, 8: 61, 2: 56, 0: 53})
Counter({6: 76, 2: 69, 5: 69, 8: 68, 4: 67, 0: 66, 1: 59, 3: 59, 9: 55, 7: 52})
Counter({8: 77, 9: 71, 4: 68, 0: 66, 2: 66, 6: 66, 7: 64, 5: 62, 1: 60, 3: 40})
Counter({8: 86, 9: 66, 4: 65, 1: 64, 2: 62, 5: 61, 0: 60, 6: 60, 3: 58, 7: 58})
Counter({7: 75, 8: 73, 6: 70, 5: 70, 3: 68, 9: 64, 4: 61, 0: 55, 2: 53, 1: 51})
Counter({6: 78, 1: 70, 5: 67, 0: 66, 2: 66, 4: 64, 8: 60, 3: 58, 9: 56, 7: 55})
Counter({9: 75, 7: 70, 8: 69, 3: 67, 4: 65, 5: 63, 2: 62, 1: 57, 0: 57, 6: 55})

@markdaoust