code: def transform_atomic_function(
      rt: runtime_client.Runtime,
      f: function_lib.AtomicFunction,
      transform_fn: Union[
          _FunctionDefTransformerType, list[_FunctionDefTransformerType]
      ],
      mlir_pipeline: Union[str, list[str]],
  ) -> function_lib.AtomicFunction:
      """Applies transforms on an AtomicFunction."""
      transform_fns = (
          transform_fn
          if isinstance(transform_fn, list)
          else [transform_fn]
      )
      mlir_pipelines = (
          mlir_pipeline
          if isinstance(mlir_pipeline, list)
          else [mlir_pipeline]
      )
      # First apply the MLIR based transformation.
      for mlir_pipeline in mlir_pipelines:
          rt.TransformFunction(
              f.cached_definition.signature.name, mlir_pipeline
          )
  
      # Get the `FunctionDef` after MLIR transformation.
      fndef = rt.GetFunctionProto(f.cached_definition.signature.name)
  
      # Apply the Python function based transformation.
      for transform_fn in transform_fns:
          transform_fn(fndef)
      rt.CreateFunction(fndef)
  
      # Generate a new `FuncGraph`
      graph = ops.get_default_graph()
      with graph.as_default():
          func_graph = function_def_lib.function_def_to_graph(
              fndef,
              structured_input_signature=f.graph.structured_input_signature,
              structured_outputs=f.graph.structured_outputs,
              propagate_device_spec=True,
          )
  
      # pylint: disable=protected-access
      # Ref: third_party/tensorflow/python/ops/control_flow_util_v2.py
      # Generate a new `AtomicFunction`.
      atomic = function_lib.from_func_graph(
          fndef.signature.name,
          func_graph,
          fndef.attr,
      )
      # pylint: enable=protected-access
  
      return atomic
  
imports: from typing import Any, Callable, Iterator, Optional, Union
  from tensorflow.core.framework import attr_value_pb2
  from tensorflow.core.framework import function_pb2
  from tensorflow.core.function.capture import restore_captures
  from tensorflow.core.function.runtime_client import runtime_client
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import function
  from tensorflow.python.framework import func_graph
  from tensorflow.python.framework import function_def_to_graph
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.ops import custom_gradient
  from tensorflow.python.ops import default_gradient
  from tensorflow.python.ops import handle_data_util
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.util import compat
label_desc: Use image node_id 2 for calling a global function with example usage: transform_atomic_function(rt, f, transform_fn, mlir_pipeline) and returns: atomic