class_docstr: Make sure that all corresponding PS gradient ops are assigned to the same task.
code: def __init__(
      self, op_to_task, agg_grad_ops, apply_grad_ops, ps_device
  ):
      self._op_to_task = op_to_task
      self._agg_grad_ops = agg_grad_ops
      self._apply_grad_ops = apply_grad_ops
      self._ps_device = ps_device
  
      # Note(gyeongin): Need to include control dependency ops in ancestors and
      # descendants or not?
      self._apply_grad_ancestor_ops = get_ancestors(
          self._apply_grad_ops, self._agg_grad_ops
      )
      self._agg_grad_descendant_ops = traverse(
          self._agg_grad_ops, end_ops=self._apply_grad_ops
      )
      self._ancestors_diff_descendants = (
          self._apply_grad_ancestor_ops.difference(
              self._agg_grad_descendant_ops
          )
      )
  
      logging.debug(
          f"apply_grad_ancestor_ops: {len(self._apply_grad_ancestor_ops)}"
      )
      logging.debug(
          f"agg_grad_descendant_ops: {len(self._agg_grad_descendant_ops)}"
      )
      logging.debug(
          f"ancestors diff descendants: {len(self._ancestors_diff_descendants)}"
      )
  
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 1 to create a new PSGradientTaskAssigner object  with example: obj = PSGradientTaskAssigner(op_to_task, agg_grad_ops, apply_grad_ops, ps_device)