code: def test_small_constants_optimization_with_grappler(self):
      def func(inp):
          x = constant_op.constant(1)
          for _ in range(1000):
              if inp:
                  x = x + constant_op.constant(1)
              else:
                  x = x + constant_op.constant(2)
          return x
  
      brancher = polymorphic_function.function(func)
      brancher_opt = polymorphic_function.function(
          func,
          experimental_attributes={
              "runtime_constant_optimization": True
          },
      )
  
      # Trace each function once.
      with ops.device_v2("CPU"):
          x = constant_op.constant(True)
      self.assertEqual(brancher(x), brancher_opt(x))
  
      benchmark = min(
          timeit.repeat(lambda: brancher(x), repeat=5, number=100)
      )
      opt_benchmark = min(
          timeit.repeat(lambda: brancher_opt(x), repeat=5, number=100)
      )
  
      # Constant folded execution is usually 15 - 20 times faster. Here we check
      # for a 2x speedup to account for various machines the test might run on.
      # Specially the kokoro machines seems to run much slower.
      self.assertLess(opt_benchmark * 2, benchmark)
  
file_path: tensorflow/tensorflow/python/eager/small_constants_optimizer_test.py
imports: import contextlib
  import timeit
  from tensorflow.python.eager import context
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the FunctionTest obj's underlying member method code with example usage: obj.test_small_constants_optimization_with_grappler() and returns: x