code: def test_strip_summary_ops(self):
      def normalize_while_node(fndef):
          """Helper method to normalize the while node for comparison."""
          for node in fndef.node_def:
              if node.op == "While":
                  # The names of the nested functions are expected to be different
                  # because they will have a uid appended to them.
                  node.attr["body"].func.name = "while_body"
                  node.attr["cond"].func.name = "while_cond"
  
                  # The summary_writer and `inlcude_summary` args are expected to be
                  # passed in and out of the transformed function as we do not modify
                  # the function signatures.
                  # Expect a mismatch in input and output types/shapes.
                  node.attr["T"].ClearField("list")
                  node.attr["output_shapes"].ClearField("list")
                  expected_inputs = {
                      "write_summary_summary_cond_input_1",
                      "record_summary",
                  }
                  if "record_summary" not in node.input:
                      continue
                  inputs = node.input
                  node.ClearField("input")
                  node.input.extend(
                      inp
                      for inp in inputs
                      if inp not in expected_inputs
                  )
                  node.attr["_num_original_outputs"].i -= 2
  
          return fndef
  
      def normalize_fdef(fndef):
          """Method to normalize the tf.function's FunctionDefs for comparison."""
          # Normalize the names for comparison as they have a uid appended.
          fndef.signature.name = "__inference_add"
  
          # The summary writer is expected to be passed into the transformed fn.
          inputs = fndef.signature.input_arg
          fndef.signature.ClearField("input_arg")
          fndef.signature.input_arg.extend(
              inp
              for inp in inputs
              if inp.name != "write_summary_summary_cond_input_1"
          )
  
          # The disable_summaries_at_runtime attr is expected to be cleared.
          fndef.attr["disable_summaries_at_runtime"].ClearField("list")
          return fndef
  
      writer = summary_ops_v2.create_file_writer_v2(self.summary_dir)
      var = variables.Variable(1.0)
  
      def remove_writer_attr(fndef):
          arg_attr = fndef.arg_attr
          attr_idx = None
          # tf.function uses TraceType to create placeholder for captures.
          # An extra "_user_specified_name" attr will be added to the placeholder.
          for idx in arg_attr:
              if (
                  arg_attr[idx].attr["_user_specified_name"].s
                  == b"input_1"
              ):
                  attr_idx = idx
                  break
          if attr_idx is not None:
              # Copy subsequent arg_attr to ensure indexes are continuous
              for idx in range(attr_idx, len(arg_attr) - 1):
                  fndef.arg_attr[idx].CopyFrom(fndef.arg_attr[idx + 1])
              del fndef.arg_attr[len(arg_attr) - 1]
          return fndef
  
      @polymorphic_function.function(
          autograph=False,
          experimental_attributes={
              "disable_summaries_at_runtime": ["record_summary", False]
          },
      )
      def add(x, y, record_summary, include_summary):
          def body(step, result):
              result += math_ops.cast(step, dtypes.float32)
              var.assign(result)
              if include_summary:
                  # Perform a summary write in a nested function.
                  with writer.as_default():
                      summary_ops_v2.set_step(step)
                      summary_ops_v2.write(
                          "my_metric", result, step=step
                      )
                  writer.flush()
              return math_ops.add(step, 1)
  
          result = math_ops.add(x, y)
          step = constant_op.constant(0, dtypes.int64)
  
          with summary_ops_v2.record_if(record_summary):
              if include_summary:
                  # Perform a summary write in the main function body.
                  with writer.as_default():
                      summary_ops_v2.set_step(step)
                      summary_ops_v2.write(
                          "my_metric", result, step=step
                      )
                  writer.flush()
  
              step = math_ops.add(step, 1)
              loop_cond = lambda i: math_ops.less(i, 3)
              loop_body = lambda i: body(i, result)
              step = while_loop.while_loop_v2(
                  loop_cond, loop_body, [step]
              )
              var.assign(result)
          return result
  
      one = constant_op.constant(1.0, dtypes.float32)
      inputs_with_summaries = [
          one,
          one,
          constant_op.constant(True),
          True,
      ]
      inputs_without_summaries = [
          one,
          one,
          constant_op.constant(False),
          False,
      ]
      inputs_without_summaries_at_runtime = [
          one,
          one,
          constant_op.constant(False),
          True,
      ]
  
      # Ensure the result of `add` is the same with and without summaries.
      self.assertEqual(
          add(*inputs_with_summaries), add(*inputs_without_summaries)
      )
  
      # Ensure the result of `add` is the same when summaries have been stripped
      # at trace time.
      self.assertEqual(
          add(*inputs_without_summaries_at_runtime),
          add(*inputs_without_summaries),
      )
  
      # Force a trace of `add` where summaries have been stripped at trace time.
      expected = add.get_concrete_function(
          *inputs_without_summaries
      ).function_def
  
      # Extract the trace of `add` where summaries have been stripped in the
      # runtime.
      function_name = add.get_concrete_function(
          *inputs_without_summaries_at_runtime
      ).function_def.signature.name
      ctx = runtime_client.GlobalPythonEagerContext()
      rt = runtime_client.Runtime(ctx)
      fndef = rt.GetFunctionProto(
          function_name + "__instance__no_summaries"
      )
  
      # Normalize the fndefs and compare them for equivalence.
      fndef = normalize_fdef(normalize_while_node(fndef))
      fndef = remove_writer_attr(fndef)
      expected = normalize_fdef(normalize_while_node(expected))
      self.assertProtoEquals(expected, fndef)
  
      # Verify that summaries were only written when executing with the
      # `inputs_with_summaries` argument.
      num_summary_events = 0
      summary_files = [
          os.path.join(self.summary_dir, sf)
          for sf in gfile.ListDirectory(self.summary_dir)
      ]
      for record in readers.TFRecordDatasetV2(
          filenames=summary_files
      ).as_numpy_iterator():
          event = event_pb2.Event()
          event.ParseFromString(record)
          if event.HasField("summary"):
              num_summary_events += 1
      # 3 Events are written by `add` when summaries are enabled.
      self.assertEqual(num_summary_events, 3)
  
file_path: tensorflow-upstream/tensorflow/python/eager/summary_optimizer_test.py
imports: import os
  import os.path
  from absl import flags
  from tensorflow.core.function.runtime_client import runtime_client
  from tensorflow.core.util import event_pb2
  from tensorflow.python.data.ops import readers
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import summary_ops_v2
  from tensorflow.python.ops import variables
  from tensorflow.python.ops import while_loop
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the SummaryOpsTransformationTest obj's underlying member method code with example usage: obj.test_strip_summary_ops() and returns: fndef, fndef, fndef, result, math_ops