class_docstr: Proxy variable implementation.
class_docstr_tok: ["Proxy","variable","implementation","."]
docstr: Update operations colocated with master variables to be colocated with proxy variables.
  
  Args:
      get_colocation_op (Callable): fn that gets the current colocation ops
docstr_tok: ["Update","operations","colocated","with","master","variables","to","be","colocated","with","proxy","variables",".","Args",":","get_colocation_op","(","Callable",")",":","fn","that","gets","the","current","colocation","ops"]
code: def update_colocation_group(self, get_colocation_op):
      """
      Update operations colocated with master variables to be colocated with proxy variables.
  
      Args:
          get_colocation_op (Callable): fn that gets the current colocation ops
      """
      for op in self._graph_item.graph.get_operations():
          # Do not update shared node (including nodes in the optimizer)
          # Do not update operations within the variable scope of master var
          # Do not update the VarhandleOp itself
          if (
              not op.name.startswith(AUTODIST_REPLICA_PREFIX)
              or op.name.startswith(self._optimizer_name_scope)
              or op.name.startswith(self._this_op.name + "/")
              or (
                  op.name.startswith(self._this_op.name)
                  and op.type == "VarHandleOp"
              )
          ):
              continue
          new_colocation_group = []
          for colocation_group in op.colocation_groups():
              current_binding_op = get_colocation_op(colocation_group)
              if current_binding_op == self._this_op:
                  op_name_to_bind_to = COLOCATION_PREFIX + as_bytes(
                      self._proxy_vars[0].op.name
                  )
                  new_colocation_group.append(op_name_to_bind_to)
              else:
                  new_colocation_group.append(colocation_group)
          op._set_attr(
              "_class",
              pb2_AttrValue(
                  list=pb2_AttrValue.ListValue(s=new_colocation_group)
              ),
          )
  
file_path: autodist/autodist/kernel/common/proxy_variable.py
imports: from collections import defaultdict
  from tensorflow.core.framework.attr_value_pb2 import AttrValue
  from tensorflow.python import ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.util.compat import as_bytes
  from autodist.const import COLOCATION_PREFIX
  from autodist.kernel.common.utils import get_consumers, update_consumers, replica_prefix, AUTODIST_REPLICA_PREFIX, parse_optimizer_scope
  from autodist.kernel.common.variable_utils import get_read_var_ops, get_read_var_tensor, gen_read_var_op
  from autodist.utils import logging
label_desc: Use image node_id 4 for calling the ProxyVariable obj's underlying member method code with example usage: obj.update_colocation_group(get_colocation_op) without return types