# Path: tensorflow_transform/py_func/pyfunc_helper.py
# _PYFUNC_COLLECTION_KEY = 'pyfuncs'
# class _PyFuncDef(tfx_namedtuple.namedtuple('_PyFuncDef', ['token', 'func'])):
#   def from_proto(attr_value, import_scope=None):
#   def from_proto_string(proto_str, import_scope=None):
#   def to_proto(self, export_scope=None):
# def insert_pyfunc(func, Tout, stateful, name, *args):  # pylint: disable=invalid-name
# def register_pyfuncs_from_saved_transform(graph, meta_graph, loaded_in_tf2):
# 
# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'
# 
# Path: tensorflow_transform/saved/saved_model_loader.py
# def parse_saved_model(saved_model_dir):
# def _choose_meta_graph_def_internal(saved_model, tags):
# def choose_meta_graph_def(saved_model):
# def choose_meta_graph_def_and_raise(saved_model):
# def get_asset_tensors(saved_model_dir, meta_graph_def_to_load):


import importlib
import os
import re
import tensorflow as tf
from tensorflow_transform.py_func import pyfunc_helper
from tensorflow_transform.saved import constants
from tensorflow_transform.saved import saved_model_loader
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.training import saver as tf_saver
        assert (name not in tensor_info_map or
                tensor_info_map[name].WhichOneof('encoding') == 'coo_sparse')
      new_tensor_info = tensor_info_map[name]
      original_tensor_type = match.group(2)
      if original_tensor_type == 'indices':
        new_tensor_info.coo_sparse.indices_tensor_name = (
            original_tensor_info.name)
      elif original_tensor_type == 'values':
        new_tensor_info.dtype = original_tensor_info.dtype
        new_tensor_info.coo_sparse.values_tensor_name = (
            original_tensor_info.name)
      elif original_tensor_type == 'dense_shape':
        new_tensor_info.coo_sparse.dense_shape_tensor_name = (
            original_tensor_info.name)
      else:
        new_tensor_info.CopyFrom(tensor_info_map[original_name])
      del tensor_info_map[original_name]


def _load_transform_saved_model(transform_savedmodel_dir):
  """Load a SavedModel representing a transform function from disk.

  Args:
    transform_savedmodel_dir: a SavedModel directory.

  Returns:
    A tuple with a `MetaGraphDef` proto, the input and outputs of a
    `SignatureDef` proto, and a dict from tensor names to absolute paths for
    asset filepaths.
  """

# Path: tensorflow_transform/py_func/pyfunc_helper.py
# _PYFUNC_COLLECTION_KEY = 'pyfuncs'
# class _PyFuncDef(tfx_namedtuple.namedtuple('_PyFuncDef', ['token', 'func'])):
#   def from_proto(attr_value, import_scope=None):
#   def from_proto_string(proto_str, import_scope=None):
#   def to_proto(self, export_scope=None):
# def insert_pyfunc(func, Tout, stateful, name, *args):  # pylint: disable=invalid-name
# def register_pyfuncs_from_saved_transform(graph, meta_graph, loaded_in_tf2):
# 
# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'
# 
# Path: tensorflow_transform/saved/saved_model_loader.py
# def parse_saved_model(saved_model_dir):
# def _choose_meta_graph_def_internal(saved_model, tags):
# def choose_meta_graph_def(saved_model):
# def choose_meta_graph_def_and_raise(saved_model):
# def get_asset_tensors(saved_model_dir, meta_graph_def_to_load):


# Path: tensorflow_transform/saved/saved_transform_io.py
import importlib
import os
import re
import tensorflow as tf
from tensorflow_transform.py_func import pyfunc_helper
from tensorflow_transform.saved import constants
from tensorflow_transform.saved import saved_model_loader
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.training import saver as tf_saver

        assert (name not in tensor_info_map or
                tensor_info_map[name].WhichOneof('encoding') == 'coo_sparse')
      new_tensor_info = tensor_info_map[name]
      original_tensor_type = match.group(2)
      if original_tensor_type == 'indices':
        new_tensor_info.coo_sparse.indices_tensor_name = (
            original_tensor_info.name)
      elif original_tensor_type == 'values':
        new_tensor_info.dtype = original_tensor_info.dtype
        new_tensor_info.coo_sparse.values_tensor_name = (
            original_tensor_info.name)
      elif original_tensor_type == 'dense_shape':
        new_tensor_info.coo_sparse.dense_shape_tensor_name = (
            original_tensor_info.name)
      else:
        new_tensor_info.CopyFrom(tensor_info_map[original_name])
      del tensor_info_map[original_name]


def _load_transform_saved_model(transform_savedmodel_dir):
  """Load a SavedModel representing a transform function from disk.

  Args:
    transform_savedmodel_dir: a SavedModel directory.

  Returns:
    A tuple with a `MetaGraphDef` proto, the input and outputs of a
    `SignatureDef` proto, and a dict from tensor names to absolute paths for
    asset filepaths.
  """

  saved_model = saved_model_loader.parse_saved_model(