def filter_dataset_on_len(dataset,
                          training,
                          len_map=None,
                          filter_on_eval=False):
  """Filters a dataset of lengths given in `len_map`.

  Args:
    dataset: `tf.data.Dataset` the dataset to filter.
    training: bool, true if we are in training mode.
    len_map: optional dict of str to (int, int). We filter examples where a
      feature's size is beyond the specified bounds. Ex:
      {'inputs': (1, 512), 'targets': (64, 128)} will keep only those examples
        where 1 <= len(inputs) <= 512 and 64 <= len(targets) <= 128.
    filter_on_eval: bool if true, we will filter in eval mode also.

  Returns:
    a filtered `tf.data.Dataset`.
  """
  if (len_map is None) or (not training and not filter_on_eval):
    return dataset

  assert isinstance(len_map, dict)
  for k, bounds in len_map.items():
    # pylint: disable=cell-var-from-loop
    # TODO(afrozm): Investigate `cell-var-from-loop` - since this is WAI and
    # there is a test too.
    def within_bounds(x, key, len_bounds):
      size = tf.shape(x[key])[0]
      min_len, max_len = len_bounds
      return (min_len <= size) and (size <= max_len)

    dataset = dataset.filter(lambda x: within_bounds(x, k, bounds))
    # pylint: enable=cell-var-from-loop

  return dataset
Filters a dataset of lengths given in `len_map`.

  Args:
    dataset: `tf.data.Dataset` the dataset to filter.
    training: bool, true if we are in training mode.
    len_map: optional dict of str to (int, int). We filter examples where a
      feature's size is beyond the specified bounds. Ex:
      {'inputs': (1, 512), 'targets': (64, 128)} will keep only those examples
        where 1 <= len(inputs) <= 512 and 64 <= len(targets) <= 128.
    filter_on_eval: bool if true, we will filter in eval mode also.

  Returns:
    a filtered `tf.data.Dataset`.
  
trax/data/tf_inputs.py