code: def pad_batch(features, batch_multiple):
    
    feature = list(features.items())[0][1]
    batch_size = tf.shape(feature)[0]
    mod = batch_size % batch_multiple
    has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)
    batch_padding = batch_multiple * has_mod - mod
  
    padded_features = {}
    for k, feature in features.items():
      rank = len(feature.shape)
      paddings = [[0, 0] for _ in range(rank)]
      paddings[0][1] = batch_padding
      padded_feature = tf.pad(feature, paddings)
      padded_features[k] = padded_feature
    return padded_features
docs: Pad batch dim of features to nearest multiple of batch_multiple.