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: Returns a new instance key for use in defining a collective op.
  
  You should call this once per each collective op of a collective instance.
  
  Args:
    group_key: the group key returned by get_group_key(). You should not
      assign the group key yourself.
    device: a canonical device string. It should be the device this collective
      op is on.
  
  Returns:
    a new instance key.
  
  Raises:
    ValueError: when the group key is invalid or the device is not in the
    group.
docstr_tok: ["Returns","a","new","instance","key","for","use","in","defining","a","collective","op",".","You","should","call","this","once","per","each","collective","op","of","a","collective","instance",".","Args",":","group_key",":","the","group","key","returned","by","get_group_key","(",")",".","You","should","not","assign","the","group","key","yourself",".","device",":","a","canonical","device","string",".","It","should","be","the","device","this","collective","op","is","on",".","Returns",":","a","new","instance","key",".","Raises",":","ValueError",":","when","the","group","key","is","invalid","or","the","device","is","not","in","the","group","."]
code: def get_instance_key(self, group_key, device):
      """Returns a new instance key for use in defining a collective op.
  
      You should call this once per each collective op of a collective instance.
  
      Args:
        group_key: the group key returned by get_group_key(). You should not
          assign the group key yourself.
        device: a canonical device string. It should be the device this collective
          op is on.
  
      Returns:
        a new instance key.
  
      Raises:
        ValueError: when the group key is invalid or the device is not in the
        group.
      """
      with self._lock:
          group = self._instance_key_table.get(group_key, None)
          if group is None:
              raise ValueError(f"Group {group_key} is not found.")
          if device not in group:
              raise ValueError(
                  f"Device {device} is not present in group {group_key}"
              )
          v = group[device]
          group[device] += 1
          return v
  
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 4 for calling the CollectiveKeys obj's underlying member method code with example usage: obj.get_instance_key(group_key, device) and returns: v