code: def test_function_jit_compile(self):
      def compute(x):
          return math_ops.reduce_mean(x, axis=0, keepdims=True)
  
      @def_function.function(jit_compile=True)
      def vectorized_compute(x):
          return pfor_control_flow_ops.vectorized_map(compute, x)
  
      result = vectorized_compute(array_ops.ones((10, 5, 3)))
      self.run_and_assert_equal(result, array_ops.ones((10, 1, 3)))
  
imports: from tensorflow.compiler.tf2xla.python import xla
  from tensorflow.python.compiler.xla import jit
  from tensorflow.python.compiler.xla import xla
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_v2_toggles
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import while_loop
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.ops.parallel_for.test_util import PForTestCase
  from tensorflow.python.platform import test
label_desc: Use image node_id 4 for calling the PForTest obj's underlying member method code with example usage: obj.test_function_jit_compile() and returns: math_ops, pfor_control_flow_ops