code: def update_colocation_group(ops, old_op, new_op):
      """
      For each op in ops, we replace the colocation group as old_op to colocation group as new_op.
  
      Args:
          ops (Iterable[Operation]): The operations to update
          old_op (Operation): The op having the old colocation group
          new_op (Operation): The op having the new colocation group
      """
      old_groups = old_op.colocation_groups() or [
          COLOCATION_PREFIX + as_bytes(new_op.name)
      ]
      new_groups = new_op.colocation_groups() or [
          COLOCATION_PREFIX + as_bytes(new_op.name)
      ]
      for op in ops:
          if op.colocation_groups() == old_groups:
              op._set_attr(
                  "_class",
                  AttrValue(list=AttrValue.ListValue(s=new_groups)),
              )
              assert op.colocation_groups() == new_groups
  
file_path: autodist/autodist/kernel/common/utils.py
imports: from collections import deque
  from tensorflow.core.framework.attr_value_pb2 import AttrValue
  from tensorflow.python.util.compat import as_bytes
  from autodist.const import AUTODIST_REPLICA_PREFIX, COLOCATION_PREFIX
label_desc: Use image node_id 12 for calling a global function with example usage: update_colocation_group(ops, old_op, new_op) without return types