def pad_fn(features: Dict[str, tf.Tensor], bsz: int) -> Dict[str, tf.Tensor]:
  """Pads a batch up to specified batch size and adds sample weights."""
  new_features = {}
  for feature_name, feature in features.items():
    actual_bsz = tf.shape(feature)[0]
    pad_size = bsz - actual_bsz
    paddings = [[0, pad_size], [0, 0]]
    new_feature = tf.pad(feature, paddings)
    new_features[feature_name] = new_feature

  sample_weights = tf.ones(actual_bsz, dtype=tf.int64)
  sample_weights = tf.pad(sample_weights, [[0, pad_size]])
  new_features['sample_weights'] = sample_weights

  return new_features
Pads a batch up to specified batch size and adds sample weights.
language/mentionmemory/utils/data_utils.py