def make_decode_fn(
    name_to_features: Dict[str, Any],
    samples_per_example: int) -> Callable[[bytes], Dict[str, tf.Tensor]]:
  """Make function to decode TF examples.

  Args:
    name_to_features: features to decode from a record in binary format.
    samples_per_example: how many TF Examples are packed in a single TF Record.

  Returns:
    A decode_fn function, which given a TF Record in binary format, parses it
      and returns features in the form of a dictionary
  """

  name_to_features = update_name_to_features(name_to_features,
                                             samples_per_example)

  def decode_fn(record: bytes) -> Dict[str, tf.Tensor]:
    """Decodes tf examples."""
    example = tf.io.parse_single_example(record, name_to_features)
    for feature, tensor in example.items():
      new_tensor_shape = (samples_per_example, tensor.shape[0] //
                          samples_per_example) + tensor.shape[1:]
      example[feature] = tf.reshape(tensor, new_tensor_shape)
    return example

  return decode_fn
Make function to decode TF examples.

  Args:
    name_to_features: features to decode from a record in binary format.
    samples_per_example: how many TF Examples are packed in a single TF Record.

  Returns:
    A decode_fn function, which given a TF Record in binary format, parses it
      and returns features in the form of a dictionary
  
language/mentionmemory/utils/data_utils.py