import tensorflow as tf


  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


@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def truncate_dataset_on_len(dataset,
                            training,
