# Path: tensorflow_transform/tf2_utils.py
# def use_tf_compat_v1(force_tf_compat_v1: bool) -> bool:
# def strip_and_get_tensors_and_control_dependencies(
#     flat_tensor_list: Iterable[tf.Tensor]
# ) -> Tuple[Iterable[tf.Tensor], Iterable[tf.Operation]]:
# def supply_missing_tensor(batch_size: int, tensor_shape: tf.TensorShape,
#                           tensor_dtype: tf.DType) -> tf.Tensor:
# def supply_missing_inputs(
#     structured_inputs: Mapping[str, common_types.TensorType],
#     batch_size: int,
#     missing_keys: Optional[Collection[str]] = None
# ) -> Mapping[str, common_types.TensorType]:
# def get_structured_inputs_from_func_graph(
#     func_graph: FuncGraph) -> Mapping[str, common_types.TensorType]:
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


import itertools
import tensorflow as tf
from tensorflow_transform import tf2_utils
from tensorflow_transform import test_case
    @tf.function(input_signature=[tf.TensorSpec([], dtype=tf.int64)])
    def func(x):
      with tf.init_scope():
        initializer_1 = tf.lookup.KeyValueTensorInitializer(
            [0, 1, 2], ['a', 'b', 'c'],
            key_dtype=tf.int64,
            value_dtype=tf.string)
        table_1 = tf.lookup.StaticHashTable(initializer_1, default_value='NAN')
        size = table_1.size()
        initializer_2 = tf.lookup.KeyValueTensorInitializer(
            ['a', 'b', 'c'], [-1, 0, 1],
            key_dtype=tf.string,
            value_dtype=tf.int64)
        table_2 = tf.lookup.StaticHashTable(initializer_2, default_value=-777)
      y = table_1.lookup(x)
      _ = table_2.lookup(y)
      z = x + size
      return {'x': x, 'z': z}

    concrete_function = func.get_concrete_function()
    flat_outputs = tf.nest.flatten(
        concrete_function.structured_outputs, expand_composites=True)
    expected_flat_outputs = [t.op.inputs[0] for t in flat_outputs]
    expected_control_dependencies = itertools.chain(
        *[t.op.control_inputs for t in flat_outputs])
    new_flat_outputs, control_dependencies = (
        tf2_utils.strip_and_get_tensors_and_control_dependencies(flat_outputs))
    self.assertEqual(new_flat_outputs, expected_flat_outputs)
    self.assertEqual(control_dependencies, set(expected_control_dependencies))


# Path: tensorflow_transform/tf2_utils.py
# def use_tf_compat_v1(force_tf_compat_v1: bool) -> bool:
# def strip_and_get_tensors_and_control_dependencies(
#     flat_tensor_list: Iterable[tf.Tensor]
# ) -> Tuple[Iterable[tf.Tensor], Iterable[tf.Operation]]:
# def supply_missing_tensor(batch_size: int, tensor_shape: tf.TensorShape,
#                           tensor_dtype: tf.DType) -> tf.Tensor:
# def supply_missing_inputs(
#     structured_inputs: Mapping[str, common_types.TensorType],
#     batch_size: int,
#     missing_keys: Optional[Collection[str]] = None
# ) -> Mapping[str, common_types.TensorType]:
# def get_structured_inputs_from_func_graph(
#     func_graph: FuncGraph) -> Mapping[str, common_types.TensorType]:
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


# Path: tensorflow_transform/tf2_utils_test.py
import itertools
import tensorflow as tf
from tensorflow_transform import tf2_utils
from tensorflow_transform import test_case

    @tf.function(input_signature=[tf.TensorSpec([], dtype=tf.int64)])
    def func(x):
      with tf.init_scope():
        initializer_1 = tf.lookup.KeyValueTensorInitializer(
            [0, 1, 2], ['a', 'b', 'c'],
            key_dtype=tf.int64,
            value_dtype=tf.string)
        table_1 = tf.lookup.StaticHashTable(initializer_1, default_value='NAN')
        size = table_1.size()
        initializer_2 = tf.lookup.KeyValueTensorInitializer(
            ['a', 'b', 'c'], [-1, 0, 1],
            key_dtype=tf.string,
            value_dtype=tf.int64)
        table_2 = tf.lookup.StaticHashTable(initializer_2, default_value=-777)
      y = table_1.lookup(x)
      _ = table_2.lookup(y)
      z = x + size
      return {'x': x, 'z': z}

    concrete_function = func.get_concrete_function()
    flat_outputs = tf.nest.flatten(
        concrete_function.structured_outputs, expand_composites=True)
    expected_flat_outputs = [t.op.inputs[0] for t in flat_outputs]
    expected_control_dependencies = itertools.chain(
        *[t.op.control_inputs for t in flat_outputs])
    new_flat_outputs, control_dependencies = (
        tf2_utils.strip_and_get_tensors_and_control_dependencies(flat_outputs))
    self.assertEqual(new_flat_outputs, expected_flat_outputs)
    self.assertEqual(control_dependencies, set(expected_control_dependencies))


  @test_case.parameters(*test_case.cross_parameters(