code: def on_device_training_loop(func):
      # Value for this attribute is from xla.DebugOptions.StepMarkerLocation.
      setattr(
          func,
          "step_marker_location",
          "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP",
      )
      return func
  
file_path: tensorflow-upstream/tensorflow/python/tpu/tpu_function.py
imports: import contextlib
  import threading
label_desc: Use image node_id 3 for calling a global function with example usage: on_device_training_loop(func) and returns: func