# Path: tensorflow_transform/pretrained_models.py
# def _get_variables(scope=None,
#                    suffix=None,
#                    collection=tf.compat.v1.GraphKeys.GLOBAL_VARIABLES):
# def _get_variables_to_restore(include=None, exclude=None):
# def apply_saved_model(model_dir, inputs, tags, signature_name=None,
#                       output_keys_in_signature=None):
# def apply_function_with_checkpoint(fn, inputs, checkpoint, include=None,
#                                    exclude=None):


import os
import tensorflow as tf
from tensorflow_transform import pretrained_models
        input_tensor_1 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_1')
        input_tensor_2 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_2')
        input_tensor_3 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_3')
        inputs = {
            'input_name1': input_tensor_1,
            'input_name2': input_tensor_2,
            'input_name3': input_tensor_3
        }
        output_tensor = pretrained_models.apply_saved_model(
            export_dir,
            inputs, [tf.saved_model.SERVING],
            signature_name='my_signature_multi_input')
        feed_dict = {input_tensor_1: [2, 3, 4, 5, 6],
                     input_tensor_2: [1, 1, 1, 1, 1],
                     input_tensor_3: [1, 1, 1, 1, -1]}
        output_value = sess.run(output_tensor, feed_dict=feed_dict)
        self.assertAllEqual(output_value, [5, 7, 9, 11, 11])

  def testApplyFunctionWithCheckpointTwoInputs(self):
    checkpoint = os.path.join(self.get_temp_dir(), 'checkpoint_two')
    self.save_checkpoint_with_two_inputs(checkpoint)
    with tf.compat.v1.Graph().as_default() as graph:
      with self.test_session(graph=graph) as sess:
        input1 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='input1')
        input2 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='input2')

# Path: tensorflow_transform/pretrained_models.py
# def _get_variables(scope=None,
#                    suffix=None,
#                    collection=tf.compat.v1.GraphKeys.GLOBAL_VARIABLES):
# def _get_variables_to_restore(include=None, exclude=None):
# def apply_saved_model(model_dir, inputs, tags, signature_name=None,
#                       output_keys_in_signature=None):
# def apply_function_with_checkpoint(fn, inputs, checkpoint, include=None,
#                                    exclude=None):


# Path: tensorflow_transform/pretrained_models_test.py
import os
import tensorflow as tf
from tensorflow_transform import pretrained_models

        input_tensor_1 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_1')
        input_tensor_2 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_2')
        input_tensor_3 = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor_3')
        inputs = {
            'input_name1': input_tensor_1,
            'input_name2': input_tensor_2,
            'input_name3': input_tensor_3
        }
        output_tensor = pretrained_models.apply_saved_model(
            export_dir,
            inputs, [tf.saved_model.SERVING],
            signature_name='my_signature_multi_input')
        feed_dict = {input_tensor_1: [2, 3, 4, 5, 6],
                     input_tensor_2: [1, 1, 1, 1, 1],
                     input_tensor_3: [1, 1, 1, 1, -1]}
        output_value = sess.run(output_tensor, feed_dict=feed_dict)
        self.assertAllEqual(output_value, [5, 7, 9, 11, 11])

  def testApplyFunctionWithCheckpointTwoInputs(self):
    checkpoint = os.path.join(self.get_temp_dir(), 'checkpoint_two')
    self.save_checkpoint_with_two_inputs(checkpoint)
    with tf.compat.v1.Graph().as_default() as graph:
      with self.test_session(graph=graph) as sess:
        input1 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='input1')
        input2 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='input2')

        output_tensor = pretrained_models.apply_function_with_checkpoint(