code: def _test_fn(self, fn, inputs):
      with ops.device("device:{}:0".format(self.device)):
          # Test with TF32 disabled
          config.enable_tensor_float_32_execution(False)
          compiled_fn = def_function.function(fn, jit_compile=True)
          hlo_text = compiled_fn.experimental_get_compiler_ir(*inputs)(
              stage="hlo"
          )
          self.assertIn("operand_precision={highest,highest}", hlo_text)
  
          # Test the output is sufficiently precise by comparing with FP64 results
          out = compiled_fn(*inputs)
          sys_details = sysconfig.get_build_info()
          if sys_details[
              "is_rocm_build"
          ]:  # MIOpen does not support fp64 data type
              f32_out = compiled_fn(
                  *[math_ops.cast(x, "float32") for x in inputs]
              )
              self.assertAllClose(out, f32_out, rtol=1e-5, atol=1e-5)
          else:
              f64_out = compiled_fn(
                  *[math_ops.cast(x, "float64") for x in inputs]
              )
              self.assertAllClose(out, f64_out, rtol=1e-5, atol=1e-5)
  
          # Test with TF32 enabled. Recompile fn because enabling TF32 does not
          # reset function cache.
          config.enable_tensor_float_32_execution(True)
          compiled_fn = def_function.function(fn, jit_compile=True)
          hlo_text = compiled_fn.experimental_get_compiler_ir(*inputs)(
              stage="hlo"
          )
          # operand_precision is not in HLO if it's the default value.
          self.assertNotIn("operand_precision", hlo_text)
  
          # On Ampere GPUs and above, which support TF32, test TF32 is used by
          # asserting outputs are not close to FP64 result
          if test_util.is_gpu_available(
              min_cuda_compute_capability=(8, 0)
          ):
              out = compiled_fn(*inputs)
              f64_out = compiled_fn(
                  *[math_ops.cast(x, "float64") for x in inputs]
              )
              self.assertNotAllClose(out, f64_out, rtol=1e-5, atol=1e-5)
  
file_path: tensorflow-upstream/tensorflow/compiler/tests/tensor_float_32_test.py
imports: from tensorflow.compiler.tests import xla_test
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.platform import googletest
  from tensorflow.python.platform import sysconfig
label_desc: Use image node_id 2 for calling the TensorFloat32Test obj's underlying member method code with example usage: obj._test_fn(fn, inputs) without return types