code: def test_tf_function_works_with_cache(self, use_tf_compat_v1):
      if not use_tf_compat_v1:
          tft_unit.skip_if_not_tf2("Tensorflow 2.x required.")
  
      def preprocessing_fn(inputs, should_add_one):
          @tf.function
          def identity(x):
              if should_add_one:
                  x = x + 1
              return x
  
          return {
              "x_mean": tft.mean(identity(inputs["x"]), name="x")
              + tf.zeros_like(inputs["x"])
          }
  
      feature_spec = {"x": tf.io.FixedLenFeature([], tf.float32)}
      input_data_dict = {
          analyzer_cache.DatasetKey("span-0"): [dict(x=-2), dict(x=4)]
      }
      run_result = self._run_pipeline(
          feature_spec,
          input_data_dict,
          functools.partial(preprocessing_fn, should_add_one=False),
          transform_fn_output_dir=None,
          use_tf_compat_v1=use_tf_compat_v1,
      )
      first_cache_output, metrics = (
          run_result.cache_output,
          run_result.metrics,
      )
  
      for key in input_data_dict:
          self.assertIn(key, first_cache_output)
          self.assertEqual(1, len(first_cache_output[key].cache_dict))
  
      self.assertMetricsCounterEqual(metrics, "num_instances", 2)
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_decoded", 0
      )
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_encoded", 1
      )
      self.assertMetricsCounterEqual(
          metrics,
          "saved_models_created",
          _SINGLE_PHASE_NUM_SAVED_MODELS,
      )
      self.assertMetricsCounterEqual(
          metrics, "analysis_input_bytes_from_cache", 0
      )
  
      # Cache is still valid since the contents of the tf.function are the same.
      run_result = self._run_pipeline(
          feature_spec,
          input_data_dict,
          functools.partial(preprocessing_fn, should_add_one=False),
          should_read_cache=True,
          transform_fn_output_dir=None,
          use_tf_compat_v1=use_tf_compat_v1,
      )
      second_cache_output, metrics = (
          run_result.cache_output,
          run_result.metrics,
      )
  
      self.assertFalse(second_cache_output)
  
      self.assertMetricsCounterEqual(metrics, "num_instances", 0)
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_decoded", 1
      )
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_encoded", 0
      )
      self.assertMetricsCounterEqual(
          metrics, "saved_models_created", _ZERO_PHASE_NUM_SAVED_MODELS
      )
      self.assertMetricsCounterGreater(
          metrics, "analysis_input_bytes_from_cache", 23
      )
  
      # Modifying the tf.function contents causes cache invalidation.
      run_result = self._run_pipeline(
          feature_spec,
          input_data_dict,
          functools.partial(preprocessing_fn, should_add_one=True),
          should_read_cache=True,
          transform_fn_output_dir=None,
          use_tf_compat_v1=use_tf_compat_v1,
      )
      third_output_cache, metrics = (
          run_result.cache_output,
          run_result.metrics,
      )
  
      for key in input_data_dict:
          self.assertIn(key, third_output_cache)
          self.assertEqual(1, len(third_output_cache[key].cache_dict))
  
      self.assertMetricsCounterEqual(metrics, "num_instances", 2)
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_decoded", 0
      )
      self.assertMetricsCounterEqual(
          metrics, "cache_entries_encoded", 1
      )
      self.assertMetricsCounterEqual(
          metrics,
          "saved_models_created",
          _SINGLE_PHASE_NUM_SAVED_MODELS,
      )
      self.assertMetricsCounterEqual(
          metrics, "analysis_input_bytes_from_cache", 0
      )
  
imports: import functools
  import os
  import struct
  from typing import Callable, Mapping, List
  import uuid
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  import tensorflow_transform
  from tensorflow_transform import analyzer_nodes
  from tensorflow_transform import common_types
  from tensorflow_transform import impl_helper
  from tensorflow_transform import nodes
  import tensorflow_transform.beam
  from tensorflow_transform.beam import analysis_graph_builder
  from tensorflow_transform.beam import analyzer_cache
  from tensorflow_transform.beam import tft_unit
  from tensorflow_transform.tf_metadata import dataset_metadata
  from tfx_bsl.types import tfx_namedtuple
label_desc: Use image node_id 14 for calling the CachedImplTest obj's underlying member method code with example usage: obj.test_tf_function_works_with_cache(use_tf_compat_v1) and returns: dict, x