code: def rank_zero_fn(
      fn: Callable[..., TReturn]
  ) -> Callable[..., Optional[TReturn]]:
      """Function that can be used as a decorator to enable a function to be called on global rank 0 only.
  
      Note:
          This decorator should be used judiciously. it should never be used on functions that need synchronization.
          It should be used very carefully with functions that mutate local state as well
  
      Example:
  
          >>> from torchtnt.utilities.distributed import rank_zero_fn
          >>> @rank_zero_fn
          ... def foo():
          ...     return 1
          ...
          >>> x = foo() # x is 1 if global rank is 0 else x is None
  
      Args:
          fn: the desired function to be executed on rank 0 only
  
      Return:
          wrapped_fn: the wrapped function that executes only if the global rank is  0
  
      """
  
      @wraps(fn)
      def wrapped_fn(*args: Any, **kwargs: Any) -> Optional[TReturn]:
          if get_global_rank() == 0:
              return fn(*args, **kwargs)
          return None
  
      return wrapped_fn
  
imports: import os
  import tempfile
  from functools import wraps
  from typing import Any, Callable, cast, List, Optional, TypeVar, Union
  import torch
  import torch.distributed
  import torch.nn.functional
  from torch import Tensor
  from torch.distributed.elastic.utils.distributed import get_free_port
  from typing_extensions import Literal
label_desc: Use image node_id 13 for calling a global function with example usage: rank_zero_fn(fn) and returns: wrapped_fn, None, fn