class_docstr: Make sure that all corresponding PS gradient ops are assigned to the same task.
docstr: Get children of and assign a task for `curr_op`. To be used as the `neighbors_fn` for `traverse`.
docstr_tok: ["Get","children","of","and","assign","a","task","for","`","curr_op","`",".","To","be","used","as","the","`","neighbors_fn","`","for","`","traverse","`","."]
code: def __assign_forward(self, curr_op, end_ops=None):
      """Get children of and assign a task for `curr_op`. To be used as the `neighbors_fn` for `traverse`."""
      end_ops = end_ops or set()
      if curr_op in self._op_to_task and curr_op not in end_ops:
          return [
              consumer
              for consumer in get_consumers(curr_op)
              if consumer in self._apply_grad_ancestor_ops
          ]
  
      placement_reference_ops = {
          input_tensor.op for input_tensor in curr_op.inputs
      }.difference(self._ancestors_diff_descendants)
  
      if not all(
          ref_op in self._op_to_task
          for ref_op in placement_reference_ops
      ):
          # At least one of `placement_reference_ops` doesn't have a task assigned yet,
          # so re-add `curr_op` to the queue and wait for them to all have tasks
          return [curr_op]
  
      self.__assign_task(curr_op, placement_reference_ops)
  
      if curr_op not in end_ops:
          return [
              consumer
              for consumer in get_consumers(curr_op)
              if consumer in self._apply_grad_ancestor_ops
          ]
  
      return []
  
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 3 for calling the PSGradientTaskAssigner obj's underlying member method code with example usage: obj.__assign_forward(curr_op, end_ops) and returns: list, unknown, list, unknown