def length_mapper(keys_to_map, dims_to_map=None, suffix="_len"):
  """Creates a mapping function to augment a tf.data.Dataset with lengths.

  Suppose we have a `dataset` with outputs `str` with shape [sequence_length].
  Here is an example usage of this function:

  dataset = dataset.map(length_mapper(['str']))

  Now the dataset will also include outputs `str_len` that can be used for
  masking out padding in a model. The 'str_len' feature will be a scalar int32
  Tensor.

  Args:
    keys_to_map: List of strings that are keys for tf.string Tensors to map to
        lengths.
    dims_to_map: List of dimensions that representing the desired length to
        retrieve. Defaults to the zero'th dimension if None.
    suffix: String to append to the given keys to indicate the mapped Tensors.

  Returns:
    _mapper: A mapping function that can be used with the tf.data.Dataset API.
  """
  if dims_to_map is None:
    dims_to_map = [0 for _ in keys_to_map]
  def _mapper(dataset):
    for k, d in zip(keys_to_map, dims_to_map):
      dataset[k + suffix] = tf.shape(dataset[k])[d]
    return dataset
  return _mapper
Creates a mapping function to augment a tf.data.Dataset with lengths.

  Suppose we have a `dataset` with outputs `str` with shape [sequence_length].
  Here is an example usage of this function:

  dataset = dataset.map(length_mapper(['str']))

  Now the dataset will also include outputs `str_len` that can be used for
  masking out padding in a model. The 'str_len' feature will be a scalar int32
  Tensor.

  Args:
    keys_to_map: List of strings that are keys for tf.string Tensors to map to
        lengths.
    dims_to_map: List of dimensions that representing the desired length to
        retrieve. Defaults to the zero'th dimension if None.
    suffix: String to append to the given keys to indicate the mapped Tensors.

  Returns:
    _mapper: A mapping function that can be used with the tf.data.Dataset API.
  
language/common/inputs/dataset_utils.py