Initial version of new tf.data.Dataset sampler.

…tobuf > 3.20.x

…ersion to 3.20.1

* Also move class_list filter to before the group_by function
* Apply the total_examples_per_class as a take() function on each
  grouped dataset
* Remove as much casting as possible from the dataset. Certain functions
  expect an int64 though and require casting.

…function.

This change means the aug and load functions apply per example now. This
will make it easier to apply random augmentations per example and is
more consistent with how we implemented it in the existing memory
sampler.

This change also improves the tests for all parts of the module.

This change adds support for passing a callable to parse the correct
class id element for batch sampling. By default y is assumed to be a 1D
tensor with the class ids and the function is lambda y:y. Otherwise we
accept an int or str and construct a parser to get the class id tensor.