def children(input_op):
  """Find successor ops for input_op.

  Args:
    input_op(tf.Operation): Input op.

  Returns: Collection of successors with a direct edge from the input_op.

  """
  return set(child_op for out in input_op.outputs
             for child_op in out.consumers())