class_docstr: This class serves as an `input_fn` for the `tf.estimator` framework.
class_docstr_tok: ["This","class","serves","as","an","`","input_fn","`","for","the","`","tf.estimator","`","framework","."]
docstr: Interface to get a data batch, fulfilling `input_fn` contract.
  
  Args:
    params: a dict containing an integer value for key 'batch_size'.
  
  Returns:
    the tuple (features, labels), where:
      - features is a dict of Tensor-valued input features; keys populated
        are:
          'image'
          'variant'
          'alt_allele_indices'
        and, if not PREDICT mode, also:
          'locus'
  
        Aside from 'image', these may be encoded specially for TPU.
  
      - label is the Tensor-valued prediction label; in train/eval
        mode the label value is is populated from the data source; in
        inference mode, the value is a constant empty Tensor value "()".
docstr_tok: ["Interface","to","get","a","data","batch",",","fulfilling","`","input_fn","`","contract",".","Args",":","params",":","a","dict","containing","an","integer","value","for","key","'batch_size","'",".","Returns",":","the","tuple","(","features",",","labels",")",",","where",":","-","features","is","a","dict","of","Tensor-valued","input","features",";","keys","populated","are",":","'image'","'variant'","'alt_allele_indices'","and",",","if","not","PREDICT","mode",",","also",":","'locus'","Aside","from","'image","'",",","these","may","be","encoded","specially","for","TPU",".","-","label","is","the","Tensor-valued","prediction","label",";","in","train\/eval","mode","the","label","value","is","is","populated","from","the","data","source",";","in","inference","mode",",","the","value","is","a","constant","empty","Tensor","value","``","(",")","''","."]
code: def __call__(self, params):
      """Interface to get a data batch, fulfilling `input_fn` contract.
  
      Args:
        params: a dict containing an integer value for key 'batch_size'.
  
      Returns:
        the tuple (features, labels), where:
          - features is a dict of Tensor-valued input features; keys populated
            are:
              'image'
              'variant'
              'alt_allele_indices'
            and, if not PREDICT mode, also:
              'locus'
  
            Aside from 'image', these may be encoded specially for TPU.
  
          - label is the Tensor-valued prediction label; in train/eval
            mode the label value is is populated from the data source; in
            inference mode, the value is a constant empty Tensor value "()".
      """
      # See https://cloud.google.com/tpu/docs/tutorials/inception-v3-advanced
      # for some background on tuning this on TPU.
  
      # TPU optimized implementation for prediction mode
      if self.mode == tf_estimator.ModeKeys.PREDICT:
          return self.prediction_input_fn(params)
  
      # Optimized following:
      #   https://www.tensorflow.org/guide/performance/datasets
      # using the information available from xprof.
      def load_dataset(filename):
          dataset = tf.data.TFRecordDataset(
              filename,
              buffer_size=self.prefetch_dataset_buffer_size,
              compression_type=compression_type,
          )
          return dataset
  
      batch_size = params["batch_size"]
      compression_type = dv_utils.compression_type_of_files(
          self.input_files
      )
  
      # NOTE: The order of the file names returned can be non-deterministic,
      # even if shuffle is false.  See internal and the note in internal.
      # We need the shuffle flag to be able to disable reordering in EVAL mode.
      dataset = tf.data.Dataset.list_files(
          [
              sharded_file_utils.normalize_to_sharded_file_pattern(
                  pattern
              )
              for pattern in self.input_file_spec.split(",")
          ],
          shuffle=self.mode == tf_estimator.ModeKeys.TRAIN,
      )
  
      # This shuffle applies to the set of files.
      # TODO: why would we shuffle the files?
      if (
          self.mode == tf_estimator.ModeKeys.TRAIN
          and self.initial_shuffle_buffer_size > 0
      ):
          dataset = dataset.shuffle(self.initial_shuffle_buffer_size)
  
      # For both TRAIN and EVAL, use the following to speed up.
      if self.sloppy:
          options = tf.data.Options()
          options.experimental_deterministic = False
          dataset = dataset.with_options(options)
      dataset = dataset.interleave(
          load_dataset,
          cycle_length=self.input_read_threads,
          num_parallel_calls=tf.data.AUTOTUNE,
      )
  
      if self.max_examples is not None:
          dataset = dataset.take(self.max_examples)
  
      if self.mode == tf_estimator.ModeKeys.TRAIN:
          dataset = dataset.repeat()
  
      # This shuffle applies to the set of records.
      if self.mode == tf_estimator.ModeKeys.TRAIN:
          if self.shuffle_buffer_size > 0:
              dataset = dataset.shuffle(self.shuffle_buffer_size)
  
      dataset = dataset.map(
          map_func=self.parse_tfexample,
          num_parallel_calls=tf.data.AUTOTUNE,
      )
      dataset = dataset.batch(
          batch_size=batch_size, drop_remainder=True
      )
      dataset = dataset.prefetch(tf.data.AUTOTUNE)
  
      return dataset
  
imports: import itertools
  from typing import Callable, Optional, Tuple
  from absl import logging
  import ml_collections
  import tensorflow
  from tensorflow import estimator
  from deepvariant import dv_constants
  from deepvariant import dv_utils
  from deepvariant import keras_modeling
  from deepvariant.protos import deepvariant_pb2
  from google.protobuf import text_format
  from third_party.nucleus.io import sharded_file_utils
label_desc: Use image node_id 4 for calling the DeepVariantInput obj's underlying member method code with example usage: obj.__call__(params) and returns: dataset, self, dataset