class_docstr: Make sure that all corresponding PS gradient ops are assigned to the same task.
docstr: Given an op, assign it a task based on the task assignments of its reference ops.
docstr_tok: ["Given","an","op",",","assign","it","a","task","based","on","the","task","assignments","of","its","reference","ops","."]
code: def __assign_task(self, curr_op, placement_reference_ops):
      """Given an op, assign it a task based on the task assignments of its reference ops."""
      placement_reference_tasks = [
          self._op_to_task[ref_op] for ref_op in placement_reference_ops
      ]
      unique_tasks = set(placement_reference_tasks)
  
      if not unique_tasks:
          raise RuntimeError(
              f"Should have placement reference for operation {curr_op.name}"
          )
  
      if len(unique_tasks) == 1:
          curr_op_task = unique_tasks.pop()
          self._op_to_task[curr_op] = curr_op_task
      else:
          # priority: assigned placement > shared
          if self.SHARED_TASK_ID in unique_tasks:
              unique_tasks.remove(self.SHARED_TASK_ID)
          if len(unique_tasks) == 1:
              curr_op_task = unique_tasks.pop()
              self._op_to_task[curr_op] = curr_op_task
          else:
              # multiple device placement -> shared
              assert len(unique_tasks) > 1
              curr_op_task = self.SHARED_TASK_ID
              self._op_to_task[curr_op] = self.SHARED_TASK_ID
  
      logging.debug(
          f"post_grad_agg_op {curr_op.name} is assigned to ps task {curr_op_task}"
      )
      if curr_op_task == self.SHARED_TASK_ID:
          # TODO: do not assign all shared ops to task 0
          # - we can do better
          curr_op_task = 0
  
      ps_device = self._ps_device.replace(task=curr_op_task)
      curr_op._set_device(ps_device)
  
imports: from functools import partial
  from typing import List
  from tensorflow.python import ops
  from tensorflow.python.framework import device_spec, dtypes, constant_op
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.ops import math_ops, data_flow_ops, gen_control_flow_ops, control_flow_ops, gen_math_ops, gen_array_ops
  from autodist.const import MAX_INT64, AUTODIST_PREFIX
  from autodist.kernel.common import utils, variable_utils
  from autodist.kernel.common.op_info import UPDATE_OP_VAR_POS
  from autodist.kernel.common.proxy_variable import ProxyVariable
  from autodist.kernel.common.utils import get_op_name, get_consumers, get_ancestors, traverse, update_consumers, update_control_consumers, replica_prefix, strip_replica_prefix, get_control_consumers, remove_from_control_consumers, get_index_from_tensor_name, update_colocation_group
  from autodist.kernel.common.variable_utils import get_read_var_ops
  from autodist.kernel.synchronization.synchronizer import Synchronizer
  from autodist.proto import synchronizers_pb2
label_desc: Use image node_id 5 for calling the PSGradientTaskAssigner obj's underlying member method code with example usage: obj.__assign_task(curr_op, placement_reference_ops) without return types