# 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
    def tensor_fn(input1, input2):
      initializer = tf.compat.v1.initializers.constant([1, 2, 3])
      with tf.compat.v1.variable_scope(
          'Model', reuse=None, initializer=initializer):
        v1 = tf.compat.v1.get_variable('v1', [3], dtype=tf.int64)
        o1 = tf.add(v1, input1, name='myadda1')
        o = tf.subtract(o1, input2, name='myadda2')
        return o
    return tensor_fn

  def save_checkpoint_with_two_inputs(self, checkpoint_path):
    test_tensor_fn = self.make_tensor_fn_two_inputs()
    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='myinputa')
        input2 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='myinputb')
        test_tensor_fn(input1, input2)
        saver = tf.compat.v1.train.Saver()
        sess.run(tf.compat.v1.global_variables_initializer())
        saver.save(sess, checkpoint_path)

  def testApplySavedModelSingleInput(self):
    export_dir = os.path.join(self.get_temp_dir(), 'single_input')
    self.save_model_with_single_input(export_dir)
    with tf.compat.v1.Graph().as_default() as graph:
      with self.test_session(graph=graph) as sess:
        input_tensor = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor')

# 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

    def tensor_fn(input1, input2):
      initializer = tf.compat.v1.initializers.constant([1, 2, 3])
      with tf.compat.v1.variable_scope(
          'Model', reuse=None, initializer=initializer):
        v1 = tf.compat.v1.get_variable('v1', [3], dtype=tf.int64)
        o1 = tf.add(v1, input1, name='myadda1')
        o = tf.subtract(o1, input2, name='myadda2')
        return o
    return tensor_fn

  def save_checkpoint_with_two_inputs(self, checkpoint_path):
    test_tensor_fn = self.make_tensor_fn_two_inputs()
    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='myinputa')
        input2 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name='myinputb')
        test_tensor_fn(input1, input2)
        saver = tf.compat.v1.train.Saver()
        sess.run(tf.compat.v1.global_variables_initializer())
        saver.save(sess, checkpoint_path)

  def testApplySavedModelSingleInput(self):
    export_dir = os.path.join(self.get_temp_dir(), 'single_input')
    self.save_model_with_single_input(export_dir)
    with tf.compat.v1.Graph().as_default() as graph:
      with self.test_session(graph=graph) as sess:
        input_tensor = tf.compat.v1.placeholder(
            dtype=tf.int32, shape=[5], name='input_tensor')

        output_tensor = pretrained_models.apply_saved_model(