class_docstr: Reader for TFRecords using the MediaSequence format.
  
  The TFrecords already contain images and spectrograms. Spectrograms are
  extracted per second and stored with size 128x100 for each second of audio.
docstr: Initializes the instance of TFRecordDatasetFactory.
  
  Initializes a data-loader using DeepMind Video Reader (DMVR) pre-processing
  (https://github.com/deepmind/dmvr).
  TFRecords are assumed to consist of tf.SequenceExample protocol buffers in
  the MediaSequence
  (https://github.com/google/mediapipe/tree/master/mediapipe/util/sequence)
  format.
  
  Args:
    base_dir: The base directory of the TFRecords.
    tables: A dictionary mapping the subset name (train, val or test) to the
      relative path of the SSTable containing them. Follows DMVR convention.
      The values of the dictionary can either be a string or a list. If it is
      a string, it specifies all the shards in the SSTable. Example -
      "/path/to/sstable@10". If passing a list, each entry is a shard of the
      SSTable. Example - "[/path/to/sstable_shard_1_of_10, ...,
      /path/to/sstabble_shard_10_of_10]." The latter scenario is useful for
      debugging.
    num_classes: The number of classes in the dataset.
    examples_per_subset:  A dictionary mapping the subset name (train, val or
      test) to the number of examples in the dataset for that subset.
    subset: The subset of the dataset to load. Must be a key of "tables"
    modalities: Which modality to load. Currently supports 'rgb' and
      'spectrogram'
    prop_data: The proportion of the data to load. If less than 1.0, this
      proportion of the total TFRecord shards are read.
    num_groups: If specified will reshard the data according to `num_groups`.
      A `group_index` should be specified if using `num_groups`.
    group_index: Index of the shard to return after resharding. `num_groups`
      should be specified if using `group_index`. This is useful in
      distributed setting where one wants to ensure that different data is
      read by different workers.
docstr_tok: ["Initializes","the","instance","of","TFRecordDatasetFactory",".","Initializes","a","data-loader","using","DeepMind","Video","Reader","(","DMVR",")","pre-processing","(","https",":","\/\/github.com\/deepmind\/dmvr",")",".","TFRecords","are","assumed","to","consist","of","tf.SequenceExample","protocol","buffers","in","the","MediaSequence","(","https",":","\/\/github.com\/google\/mediapipe\/tree\/master\/mediapipe\/util\/sequence",")","format",".","Args",":","base_dir",":","The","base","directory","of","the","TFRecords",".","tables",":","A","dictionary","mapping","the","subset","name","(","train",",","val","or","test",")","to","the","relative","path","of","the","SSTable","containing","them",".","Follows","DMVR","convention",".","The","values","of","the","dictionary","can","either","be","a","string","or","a","list",".","If","it","is","a","string",",","it","specifies","all","the","shards","in","the","SSTable",".","Example","-","``","\/path\/to\/sstable","@","10","''",".","If","passing","a","list",",","each","entry","is","a","shard","of","the","SSTable",".","Example","-","``","[","\/path\/to\/sstable_shard_1_of_10",",","...",",","\/path\/to\/sstabble_shard_10_of_10","]",".","''","The","latter","scenario","is","useful","for","debugging",".","num_classes",":","The","number","of","classes","in","the","dataset",".","examples_per_subset",":","A","dictionary","mapping","the","subset","name","(","train",",","val","or","test",")","to","the","number","of","examples","in","the","dataset","for","that","subset",".","subset",":","The","subset","of","the","dataset","to","load",".","Must","be","a","key","of","``","tables","''","modalities",":","Which","modality","to","load",".","Currently","supports","'rgb","'","and","'spectrogram'","prop_data",":","The","proportion","of","the","data","to","load",".","If","less","than","1.0",",","this","proportion","of","the","total","TFRecord","shards","are","read",".","num_groups",":","If","specified","will","reshard","the","data","according","to","`","num_groups","`",".","A","`","group_index","`","should","be","specified","if","using","`","num_groups","`",".","group_index",":","Index","of","the","shard","to","return","after","resharding",".","`","num_groups","`","should","be","specified","if","using","`","group_index","`",".","This","is","useful","in","distributed","setting","where","one","wants","to","ensure","that","different","data","is","read","by","different","workers","."]
code: def __init__(
      self,
      base_dir: str,
      tables: Dict[str, Union[str, List[str]]],
      num_classes: int,
      examples_per_subset: Dict[str, int],
      subset: str = "train",
      modalities: Tuple[str] = ("rgb",),
      prop_data: float = 1.0,
      num_groups: Optional[int] = None,
      group_index: Optional[int] = None,
  ):
      """Initializes the instance of TFRecordDatasetFactory.
  
      Initializes a data-loader using DeepMind Video Reader (DMVR) pre-processing
      (https://github.com/deepmind/dmvr).
      TFRecords are assumed to consist of tf.SequenceExample protocol buffers in
      the MediaSequence
      (https://github.com/google/mediapipe/tree/master/mediapipe/util/sequence)
      format.
  
      Args:
        base_dir: The base directory of the TFRecords.
        tables: A dictionary mapping the subset name (train, val or test) to the
          relative path of the SSTable containing them. Follows DMVR convention.
          The values of the dictionary can either be a string or a list. If it is
          a string, it specifies all the shards in the SSTable. Example -
          "/path/to/sstable@10". If passing a list, each entry is a shard of the
          SSTable. Example - "[/path/to/sstable_shard_1_of_10, ...,
          /path/to/sstabble_shard_10_of_10]." The latter scenario is useful for
          debugging.
        num_classes: The number of classes in the dataset.
        examples_per_subset:  A dictionary mapping the subset name (train, val or
          test) to the number of examples in the dataset for that subset.
        subset: The subset of the dataset to load. Must be a key of "tables"
        modalities: Which modality to load. Currently supports 'rgb' and
          'spectrogram'
        prop_data: The proportion of the data to load. If less than 1.0, this
          proportion of the total TFRecord shards are read.
        num_groups: If specified will reshard the data according to `num_groups`.
          A `group_index` should be specified if using `num_groups`.
        group_index: Index of the shard to return after resharding. `num_groups`
          should be specified if using `group_index`. This is useful in
          distributed setting where one wants to ensure that different data is
          read by different workers.
      """
  
      for modality in modalities:
          if modality not in AVTFRecordDatasetFactory._MODALITIES:
              raise ValueError("Invalid modality %s." % modality)
      self._modalities = modalities
  
      super().__init__(
          base_dir=base_dir,
          tables=tables,
          examples_per_subset=examples_per_subset,
          subset=subset,
          num_classes=num_classes,
          fraction_data=prop_data,
          num_groups=num_groups,
          group_index=group_index,
      )
  
imports: import functools
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import modalities
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.dataset_lib import video_ops
  from scenic.projects.mbt.datasets.dataset_utils import add_spectrogram
  from scenic.projects.vivit.data import video_tfrecord_dataset
  import tensorflow
label_desc: Use image node_id 1 to create a new AVTFRecordDatasetFactory object from inherited base classes: video_tfrecord_dataset with example: obj = AVTFRecordDatasetFactory(base_dir, tables, num_classes, examples_per_subset, subset, modalities, prop_data, num_groups, group_index)