class_docstr: Class that manages collective keys.
  
  We need to manage three different keys for collective:
  
  *Group key*: an integer key to identify the set of cooperative devices.
  Collective ops work under the same set of devices must using the same group
  key.
  
  *Instance key*: an integer key to identify the set of same counterpart of
  tensors on different devices in a device group that need to be all-reduced.
  
  This class is thread safe.
class_docstr_tok: ["Class","that","manages","collective","keys",".","We","need","to","manage","three","different","keys","for","collective",":","*","Group","key","*",":","an","integer","key","to","identify","the","set","of","cooperative","devices",".","Collective","ops","work","under","the","same","set","of","devices","must","using","the","same","group","key",".","*","Instance","key","*",":","an","integer","key","to","identify","the","set","of","same","counterpart","of","tensors","on","different","devices","in","a","device","group","that","need","to","be","all-reduced",".","This","class","is","thread","safe","."]
docstr: Initializes the object.
  
  Args:
    group_key_start: the starting integer of group key.
docstr_tok: ["Initializes","the","object",".","Args",":","group_key_start",":","the","starting","integer","of","group","key","."]
code: def __init__(self, group_key_start=1):
      """Initializes the object.
  
      Args:
        group_key_start: the starting integer of group key.
      """
      self._group_key = group_key_start
      self._instance_key_table = {}
      self._lock = threading.Lock()
      self._known_groups = {}
  
imports: import copy
  import threading
  from typing import Callable, List, Optional, Union
  from tensorflow.python.distribute import collective_util
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import backprop_util
  from tensorflow.python.eager import context
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import collective_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nccl_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.types import core
label_desc: Use image node_id 1 to create a new CollectiveKeys object from inherited base classes: object with example: obj = CollectiveKeys(group_key_start)