docstr: Gets the mnist function.
  
  Args:
    use_trt: whether use TF-TRT to convert the graph.
    model_dir: the model directory to load the checkpoints.
    use_dynamic_shape: whether to run the TF-TRT conversion in dynamic shape
      mode.
  
  Returns:
    The mnist model function.
code: def _GetFunc(self, use_trt, model_dir, use_dynamic_shape):
      """Gets the mnist function.
  
      Args:
        use_trt: whether use TF-TRT to convert the graph.
        model_dir: the model directory to load the checkpoints.
        use_dynamic_shape: whether to run the TF-TRT conversion in dynamic shape
          mode.
  
      Returns:
        The mnist model function.
      """
      with tempfile.TemporaryDirectory() as tmpdir:
          saved_model_dir = os.path.join(tmpdir, "mnist")
          self._SaveModel(model_dir, saved_model_dir)
  
          if use_trt:
              conv_params = trt_convert.TrtConversionParams(
                  precision_mode="FP16",
                  minimum_segment_size=2,
                  max_workspace_size_bytes=(
                      trt_convert.DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES
                  ),
                  maximum_cached_engines=1,
              )
              converter = trt_convert.TrtGraphConverterV2(
                  input_saved_model_dir=saved_model_dir,
                  use_dynamic_shape=use_dynamic_shape,
                  dynamic_shape_profile_strategy="ImplicitBatchModeCompatible",
                  **conv_params._asdict(),
              )
              converter.convert()
              try:
                  line_length = max(160, os.get_terminal_size().columns)
              except OSError:
                  line_length = 160
              converter.summary(line_length=line_length, detailed=True)
              func = converter._converted_func
          else:
              saved_model_loaded = saved_model_load(
                  saved_model_dir, tags=[tag_constants.SERVING]
              )
              func = saved_model_loaded.signatures[
                  signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
              ]
      return func
  
imports: import os.path
  import tempfile
  import tensorflow_datasets
  from tensorflow.compiler.tf2tensorrt._pywrap_py_utils import is_tensorrt_enabled
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.client import session
  from tensorflow.python.compiler.tensorrt import trt_convert
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.estimator.estimator import Estimator
  from tensorflow.python.estimator.model_fn import EstimatorSpec
  from tensorflow.python.estimator.model_fn import ModeKeys
  from tensorflow.python.estimator.run_config import RunConfig
  from tensorflow.python.framework import convert_to_constants
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.keras.metrics import Accuracy
  from tensorflow.python.layers import layers
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import metrics
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.saved_model import builder
  from tensorflow.python.saved_model import signature_constants
  from tensorflow.python.saved_model import tag_constants
  from tensorflow.python.saved_model import utils
  from tensorflow.python.saved_model import signature_def_utils
  from tensorflow.python.saved_model.load import load
  from tensorflow.python.summary import summary
  from tensorflow.python.training import saver
  from tensorflow.python.training.adam import AdamOptimizer
  from tensorflow.python.training.checkpoint_management import latest_checkpoint
  from tensorflow.python.training.training_util import get_global_step
label_desc: Use image node_id 2 for calling the MNISTTestV2 obj's underlying member method code with example usage: obj._GetFunc(use_trt, model_dir, use_dynamic_shape) and returns: func