import tensorflow as tf

def get_sorted_full_neighbor(nodes, edge_types):
  """
  Args:
    nodes: A `Tensor` of `int64`.
    edge_types: A 1-D `Tensor` of int32. Specify edge types to filter outgoing
      edges.

  Return:
    A tuple of `SparseTensor` (neibors, weights).
      neighbors: A `SparseTensor` of `int64`.
      weights: A `SparseTensor` of `float`.
      types: A `SparseTensor` of `int32`
  """
  sp_returns = base._LIB_OP.get_sorted_full_neighbor(nodes, edge_types)
  return tf.SparseTensor(*sp_returns[:3]), tf.SparseTensor(*sp_returns[3:6]), \
         tf.SparseTensor(*sp_returns[6:])


def sample_fanout(nodes, edge_types, counts, default_node=-1):
  """
  Sample multi-hop neighbors of nodes according to weight in graph.

  Args:
    nodes: A 1-D `Tensor` of `int64`.
    edge_types: A list of 1-D `Tensor` of int32. Specify edge types to filter
      outgoing edges in each hop.
    counts: A list of `int`. Specify the number of sampling for each node in
      each hop.
