code: def _all_gather_keys(local_dict: Dict[Any, Any]) -> List[Any]:
      """Gathers all keys, and returns them sorted."""
      keys = list(local_dict.keys())
      gathered_keys: List[List[Any]] = [None] * dist.get_world_size()  # type: ignore[list-item]
  
      dist.all_gather_object(gathered_keys, keys)
      return sorted(set(itertools.chain.from_iterable(gathered_keys)))
  
imports: import cProfile
  import io
  import itertools
  import os
  from contextlib import contextmanager
  from pstats import Stats
  from typing import Any, Callable, cast, Dict, List, Optional, Sequence, TypeVar, Union
  import torch
  import torch.distributed
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._shard.sharded_tensor.shard import Shard
  from torch.distributed._tensor import DTensor
  from .api import _is_wrapped_exception, _wrap_exception, CheckpointException, WRAPPED_EXCEPTION
  from .metadata import MetadataIndex, STATE_DICT_TYPE
label_desc: Use image node_id 2 for calling a global function with example usage: _all_gather_keys(local_dict) and returns: sorted