# Path: tensorflow_transform/saved/saved_transform_io.py
# _MANGLED_TENSOR_NAME_RE = re.compile(
#     r'(.*)\$(indices|values|dense_shape|dense_tensor)$')
# _PARTITIONED_VARIABLE_NAME_RE = re.compile(r'^(.*)/part_(\d*)$')
# def _update_legacy_signature(signature):
# def _load_transform_saved_model(transform_savedmodel_dir):
# def _expand_input_map(logical_input_map, input_signature):
# def _maybe_register_addon_ops():
#   def _try_import(name):
# def _partially_apply_saved_transform_impl(saved_model_dir,
#                                           logical_input_map,
#                                           tensor_replacement_map=None):
#   def lookup_remapped_tensor(tensor_name):
#   def lookup_tensor_or_sparse_or_composite_tensor(tensor_info):
# def partially_apply_saved_transform_internal(saved_model_dir,
#                                              logical_input_map,
#                                              tensor_replacement_map=None):
# def write_saved_transform_from_session(
#     session, inputs, outputs, export_path, as_text=False):
# def exported_as_v1(transform_savedmodel_dir):


import os
import tempfile
import numpy as np
import tensorflow as tf
from tensorflow_transform.saved import saved_transform_io
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import lookup_ops
        # indices and shape unchanged; values multiplied by 10 and divided by 5
        self.assertEqual(indices.tolist(), result.indices.tolist())
        self.assertEqual([2.0, 4.0], result.values.tolist())
        self.assertEqual(shape.tolist(), result.dense_shape.tolist())

  def test_ragged_roundtrip(self):
    if not hasattr(meta_graph_pb2.TensorInfo, 'CompositeTensor'):
      self.skipTest('This version of TensorFlow does not support '
                    'CompositeTenors in TensorInfo.')
    export_path = os.path.join(tempfile.mkdtemp(), 'export')

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        input_float = tf.compat.v1.ragged.placeholder(tf.float32, ragged_rank=1,
                                                      value_shape=[])
        output = input_float / 2.0
        inputs = {'input': input_float}
        outputs = {'output': output}
        saved_transform_io.write_saved_transform_from_session(
            session, inputs, outputs, export_path)

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        splits = np.array([0, 2, 3], dtype=np.int64)
        values = np.array([1.0, 2.0, 4.0], dtype=np.float32)
        input_ragged = tf.RaggedTensor.from_row_splits(values, splits)

        # Using a computed input gives confidence that the graphs are fused
        inputs = {'input': input_ragged * 10}
        _, outputs = (

# Path: tensorflow_transform/saved/saved_transform_io.py
# _MANGLED_TENSOR_NAME_RE = re.compile(
#     r'(.*)\$(indices|values|dense_shape|dense_tensor)$')
# _PARTITIONED_VARIABLE_NAME_RE = re.compile(r'^(.*)/part_(\d*)$')
# def _update_legacy_signature(signature):
# def _load_transform_saved_model(transform_savedmodel_dir):
# def _expand_input_map(logical_input_map, input_signature):
# def _maybe_register_addon_ops():
#   def _try_import(name):
# def _partially_apply_saved_transform_impl(saved_model_dir,
#                                           logical_input_map,
#                                           tensor_replacement_map=None):
#   def lookup_remapped_tensor(tensor_name):
#   def lookup_tensor_or_sparse_or_composite_tensor(tensor_info):
# def partially_apply_saved_transform_internal(saved_model_dir,
#                                              logical_input_map,
#                                              tensor_replacement_map=None):
# def write_saved_transform_from_session(
#     session, inputs, outputs, export_path, as_text=False):
# def exported_as_v1(transform_savedmodel_dir):


# Path: tensorflow_transform/saved/saved_transform_io_test.py
import os
import tempfile
import numpy as np
import tensorflow as tf
from tensorflow_transform.saved import saved_transform_io
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import lookup_ops

        # indices and shape unchanged; values multiplied by 10 and divided by 5
        self.assertEqual(indices.tolist(), result.indices.tolist())
        self.assertEqual([2.0, 4.0], result.values.tolist())
        self.assertEqual(shape.tolist(), result.dense_shape.tolist())

  def test_ragged_roundtrip(self):
    if not hasattr(meta_graph_pb2.TensorInfo, 'CompositeTensor'):
      self.skipTest('This version of TensorFlow does not support '
                    'CompositeTenors in TensorInfo.')
    export_path = os.path.join(tempfile.mkdtemp(), 'export')

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        input_float = tf.compat.v1.ragged.placeholder(tf.float32, ragged_rank=1,
                                                      value_shape=[])
        output = input_float / 2.0
        inputs = {'input': input_float}
        outputs = {'output': output}
        saved_transform_io.write_saved_transform_from_session(
            session, inputs, outputs, export_path)

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        splits = np.array([0, 2, 3], dtype=np.int64)
        values = np.array([1.0, 2.0, 4.0], dtype=np.float32)
        input_ragged = tf.RaggedTensor.from_row_splits(values, splits)

        # Using a computed input gives confidence that the graphs are fused
        inputs = {'input': input_ragged * 10}
        _, outputs = (

            saved_transform_io.partially_apply_saved_transform_internal(