class_docstr: Make sure that all corresponding PS gradient ops are assigned to the same task.
docstr: Bi-directionally traverse the graph and assign tasks to ops.
docstr_tok: ["Bi-directionally","traverse","the","graph","and","assign","tasks","to","ops","."]
code: def assign(self):
      """Bi-directionally traverse the graph and assign tasks to ops."""
      # Parent-to-child traversal
      fn = partial(self.__assign_forward, end_ops=self._apply_grad_ops)
      traverse(
          self._agg_grad_ops,
          end_ops=self._apply_grad_ops,
          neighbors_fn=fn,
      )
      # Child-to-parent traversal
      fn = partial(self.__assign_backward, end_ops=self._agg_grad_ops)
      traverse(
          self._apply_grad_ops, self._agg_grad_ops, neighbors_fn=fn
      )
  
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 2 for calling the PSGradientTaskAssigner obj's underlying member method code with example usage: obj.assign() without return types