code: def implementation_selecting(fn_name, default_fn, cpu_fn):
      """Secret handshake: Grappler replaces default_fn w/ cpu_fn when run on CPU.
  
      Note that this should generally be called in the scope of a `tf.function`.
  
      Example:
      ```python
      @tf.function(autograph=False)
      def cumsum(x):
  
        def _xla_friendly(x):
          return tf.math.cumsum(x)
  
        def _xla_hostile(x):
          return tf.while_loop(
              cond=lambda x, _: tf.size(x) > 0,
              body=lambda x, cx: (
                  x[:-1], cx + tf.pad(x, [[tf.size(cx) - tf.size(x), 0]])),
              loop_vars=[x, tf.zeros_like(x)],
              shape_invariants=[tf.TensorShape([None]), x.shape])[1]
  
        return implementation_selecting(
            'cumsum', default_fn=_xla_friendly, cpu_fn=_xla_hostile)(x=x)
  
      retval, which_runtime = cumsum(tf.constant([1, 2, 3]))
      ```
  
      Args:
        fn_name: A `str` prefix used to specify the "API" being implemented by these
          functions. Mostly only interesting for debugging / graph inspection.
        default_fn: The default implementation. This is used when we are in an XLA
          tracing context, in TFP-on-JAX, or running on devices such as GPU or TPU.
        cpu_fn: The CPU-specialized implementation. This impl is used when Grappler
          sees that the function execution has been placed on the CPU. (Also used
          for TFP-on-numpy.)
  
      Returns:
        impl_selecting_fn: A function which takes kwargs only, forwards them to the
          underlying functions, and returns a tuple:
          `(selected_fn(**kwargs), runtime_that_was_selected)`.
      """
  
      def wrap_fn(fn, runtime):
          def stub_fn(**kwargs):
              return fn(**kwargs), tf.constant(
                  runtime, dtype=tf.float32
              )
  
          stub_fn.__name__ = fn.__name__
          return stub_fn
  
      default_fn = wrap_fn(default_fn, _RUNTIME_DEFAULT)
      cpu_fn = wrap_fn(cpu_fn, _RUNTIME_CPU)
  
      def impl_selecting_fn(**kwargs):
          """The wrapper function to be returned."""
          if is_xla():  # JAX, XLA breakout.
              return default_fn(**kwargs)
          if NUMPY_MODE:  # Numpy breakout.
              return cpu_fn(**kwargs)
  
          # Each time a `tf.function` is called, we will give it a unique
          # identifiable API name, so that Grappler won't get confused when it
          # sees multiple samplers in same graph, and it will be able
          # to pair up the different implementations across them.
          api_name = "{}_{}".format(fn_name, str(uuid.uuid4()))
          default_fn_with_backend = _function_with_backend(
              default_fn, api_name
          )
          cpu_fn_with_backend = _function_with_backend(
              cpu_fn, api_name, preferred_device=_CPU_DEVICE_NAME
          )
  
          # Call the default sampling impl and register the CPU-specialized impl.
          # Grappler will kick in during session execution to optimize the graph.
          result, runtime = default_fn_with_backend(**kwargs)
          concrete_func = cpu_fn_with_backend.get_concrete_function(
              **kwargs
          )
          concrete_func.add_to_graph()
          concrete_func.add_gradient_functions_to_graph()
          return result, runtime
  
      return impl_selecting_fn
  
imports: import functools
  import uuid
  import tensorflow.compat.v2
label_desc: Use image node_id 4 for calling a global function with example usage: implementation_selecting(fn_name, default_fn, cpu_fn) and returns: impl_selecting_fn, stub_fn, result,  runtime, fn,  tf, default_fn, cpu_fn