code: def testReduceMeanOverflow(self):
      @function(jit_compile=True)
      def tpu_computation():
          shape = (4, 32, 512, 512, 96)
          dtype = dtypes.float32
          test_ones = array_ops.ones(shape=shape, dtype=dtype)
          return math_ops.reduce_mean(test_ones)
  
      with ops.device("TPU:0"):
          result = tpu_computation()
          self.assertAllClose(result, 1.0, rtol=1e-4, atol=1e-4)
  
file_path: tensorflow/tensorflow/compiler/tests/mean_op_test.py
imports: from tensorflow.python.eager.def_function import function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the MeanOpTest obj's underlying member method code with example usage: obj.testReduceMeanOverflow() and returns: math_ops