code: def enable_argscope_for_function(func, log_shape=True):
      """Decorator for function to support argscope
  
      Example:
  
          .. code-block:: python
  
              from mylib import myfunc
              myfunc = enable_argscope_for_function(myfunc)
  
      Args:
          func: A function mapping one or multiple tensors to one or multiple
              tensors.
          log_shape (bool): Specify whether the first input resp. output tensor
              shape should be printed once.
  
      Remarks:
          If the function ``func`` returns multiple input or output tensors,
          only the first input/output tensor shape is displayed during logging.
  
      Returns:
          The decorated function.
  
      """
  
      assert callable(func), "func should be a callable"
  
      @wraps(func)
      def wrapped_func(*args, **kwargs):
          actual_args = copy.copy(get_arg_scope()[func.__name__])
          actual_args.update(kwargs)
          out_tensor = func(*args, **actual_args)
          in_tensor = args[0]
  
          ctx = get_current_tower_context()
          name = (
              func.__name__ if "name" not in kwargs else kwargs["name"]
          )
          if log_shape:
              if (
                  "tower" not in ctx.ns_name.lower()
              ) or ctx.is_main_training_tower:
                  # we assume the first parameter is the most interesting
                  if isinstance(out_tensor, tuple):
                      out_tensor_descr = out_tensor[0]
                  else:
                      out_tensor_descr = out_tensor
                  logger.info(
                      "{:<12}: {} --> {}".format(
                          "'" + name + "'",
                          get_shape_str(in_tensor),
                          get_shape_str(out_tensor_descr),
                      )
                  )
  
          return out_tensor
  
      wrapped_func.__argscope_enabled__ = True
      return wrapped_func
  
file_path: tensorpack/tensorpack/tfutils/argscope.py
imports: import copy
  from collections import defaultdict
  from contextlib import contextmanager
  from functools import wraps
  from inspect import getmembers, isfunction
  import tensorflow
  from ..compat import is_tfv2
  from ..utils import logger
  from .model_utils import get_shape_str
  from .tower import get_current_tower_context
label_desc: Use image node_id 3 for calling a global function with example usage: enable_argscope_for_function(func, log_shape) and returns: wrapped_func, out_tensor