class_docstr: Testing the output shape of a TRTEngine.
class_docstr_tok: ["Testing","the","output","shape","of","a","TRTEngine","."]
code: def _GetInferGraph(self, *args, **kwargs):
      trt_saved_model_dir = super(
          TRTEngineOpInputOutputShapeTest, self
      )._GetInferGraph(*args, **kwargs)
  
      def get_func_from_saved_model(saved_model_dir):
          try:  # Necessary for `bazel run ...`
              saved_model_load_fn = load.load
          except AttributeError:  # All the other cases
              saved_model_load_fn = load
          saved_model_loaded = saved_model_load_fn(
              saved_model_dir, tags=[tag_constants.SERVING]
          )
          graph_func = saved_model_loaded.signatures[
              signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
          ]
          return graph_func, saved_model_loaded
  
      func, _ = get_func_from_saved_model(trt_saved_model_dir)
  
      input_shape = func.inputs[0].shape
      if isinstance(input_shape, tensor_shape.TensorShape):
          input_shape = input_shape.as_list()
  
      output_shapes = [
          out_shape.shape.as_list()
          if isinstance(out_shape.shape, tensor_shape.TensorShape)
          else out_shape.shape
          for out_shape in func.outputs
      ]
  
      self.assertEqual(func.inputs[0].dtype, dtypes.float32)
      self.assertEqual(func.outputs[0].dtype, dtypes.float32)
      self.assertEqual(func.outputs[1].dtype, dtypes.float32)
  
      self.assertEqual(input_shape, [None, 2, 1, 4])
      self.assertEqual(output_shapes[0], [None, 2, 4])
      self.assertEqual(output_shapes[1], [])
  
      return trt_saved_model_dir
  
file_path: tensorflow-upstream/tensorflow/python/compiler/tensorrt/test/trt_engine_op_shape_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.platform import test
  from tensorflow.python.saved_model import load
  from tensorflow.python.saved_model import signature_constants
  from tensorflow.python.saved_model import tag_constants
label_desc: Use image node_id 3 for calling the TRTEngineOpInputOutputShapeTest obj's underlying member method code with example usage: obj._GetInferGraph() and returns: trt_saved_model_dir, graph_func,  saved_model_loaded