def get_multi_hop_neighbor(nodes, edge_types):
"""
Get multi-hop neighbors with adjacent matrix.
Args:
nodes: A 1-D `tf.Tensor` of `int64`.
edge_types: A list of 1-D `tf.Tensor` of `int32`. Specify edge types to
filter outgoing edges in each hop.
Return:
A tuple of list: (nodes, adjcents)
nodes: A list of N + 1 `tf.Tensor` of `int64`, N is the number of
hops. Specify node set of each hop, including the root.
adjcents: A list of N `tf.SparseTensor` of `int64`. Specify adjacent
matrix between hops.
"""
edge_types = [type_ops.get_edge_type_id(edge_type)
for edge_type in edge_types]
nodes = tf.reshape(nodes, [-1])
nodes_list = [nodes]
adj_list = []
for hop_edge_types in edge_types:
neighbor, weight, _ = get_full_neighbor(nodes, hop_edge_types)
next_nodes, next_idx = tf.unique(neighbor.values, out_idx=tf.int64)
next_indices = tf.stack([neighbor.indices[:, 0], next_idx], 1)
next_values = weight.values
next_shape = tf.stack([tf.size(nodes), tf.size(next_nodes)])
next_shape = tf.cast(next_shape, tf.int64)
next_adj = tf.SparseTensor(next_indices, next_values, next_shape)
next_adj = tf.sparse_reorder(next_adj)
nodes_list.append(next_nodes)
adj_list.append(next_adj)
nodes = next_nodes
return nodes_list, adj_list
def sample_fanout_layerwise_each_node(nodes, edge_types, counts,
default_node=-1):
'''
sample fanout layerwise for each node
'''
edge_types = [type_ops.get_edge_type_id(edge_type)
for edge_type in edge_types]
neighbors_list = [tf.reshape(nodes, [-1])]
adj_list = []
for hop_edge_types, count in zip(edge_types, counts):
if (len(neighbors_list) == 1):
neighbors, _, _ = sample_neighbor(neighbors_list[-1],
hop_edge_types,
count,
default_node=default_node)
neighbors_list.append(tf.reshape(neighbors, [-1]))
else:
neighbors, adj = sample_neighbor_layerwise(
tf.reshape(neighbors_list[-1], [-1, last_count]),
hop_edge_types,
count,
default_node=default_node)
neighbors_list.append(tf.reshape(neighbors, [-1]))
adj_list.append(adj)
last_count = count
return neighbors_list, adj_list
def random_walk(nodes, edge_types, p=1.0, q=1.0, default_node=-1):
'''
Random walk from a list of nodes.
Args:
nodes: start node ids, 1-d Tensor
edge_types: list of 1-d Tensor of edge types
p: back probality
q: forward probality
default_node: default fill nodes
'''
edge_types = [type_ops.get_edge_type_id(edge_type)
for edge_type in edge_types]
return _random_walk(nodes, edge_types, p, q, default_node)
def sample_edge(count, edge_type=None):
"""
Sample Edges by specific types
Args:
count: A scalar tensor specify sample count
types: A scalar tensor specify sample type
Return:
A 2-d tensor of sample edge ids
"""
if edge_type == '-1':
types = -1
else:
types = type_ops.get_edge_type_id(edge_type)
return base._LIB_OP.sample_edge(count, types)
def get_full_neighbor(nodes, edge_types, condition=''):
"""
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`
"""
edge_types = type_ops.get_edge_type_id(edge_types)
sp_returns = base._LIB_OP.get_full_neighbor(nodes, edge_types, condition)
return tf.SparseTensor(*sp_returns[:3]), \
tf.SparseTensor(*sp_returns[3:6]), \
tf.SparseTensor(*sp_returns[6:])
def sample_fanout_layerwise(nodes, edge_types, counts,
default_node=-1, weight_func=''):
edge_types = [type_ops.get_edge_type_id(edge_type)
for edge_type in edge_types]
neighbors_list = [tf.reshape(nodes, [-1])]
adj_list = []
last_count = tf.size(nodes)
for hop_edge_types, count in zip(edge_types, counts):
neighbors, adj = sample_neighbor_layerwise(
tf.reshape(neighbors_list[-1], [-1, last_count]),
hop_edge_types,
count,
default_node=default_node,
weight_func=weight_func)
neighbors_list.append(tf.reshape(neighbors, [-1]))
adj_list.append(adj)
last_count = count
return neighbors_list, adj_list
def sample_fanout_with_feature(nodes, edge_types, count, default_node,
dense_feature_names, dense_dimensions,
sparse_feature_names, sparse_default_values):
edge_types = type_ops.get_edge_type_id(edge_types)
res = _sample_fanout_with_feature(
tf.reshape(nodes, [-1]), edge_types, count,
default_node=default_node,
sparse_feature_names=sparse_feature_names,
sparse_default_values=sparse_default_values,
dense_feature_names=dense_feature_names,
dense_dimensions=dense_dimensions,
N=len(count),
ND=(len(count) + 1) * len(dense_feature_names),
NS=(len(count) + 1) * len(sparse_feature_names))
neighbors = [tf.reshape(nodes, [-1])]
neighbors.extend([tf.reshape(i, [-1]) for i in res[0]])
weights = res[1]
types = res[2]
dense_features = res[3]
sparse_features = [tf.SparseTensor(*sp) for sp in zip(*res[4:7])]
return neighbors, weights, types, dense_features, sparse_features
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.
default_node: A `int`. Specify the node id to fill when there is no
neighbor for specific nodes.
Return:
A tuple of list: (samples, weights)
samples: A list of `Tensor` of `int64`, with the same length as
`edge_types` and `counts`, with shapes `[num_nodes]`,
`[num_nodes * count1]`, `[num_nodes * count1 * count2]`, ...
weights: A list of `Tensor` of `float`, with shapes
`[num_nodes * count1]`, `[num_nodes * count1 * count2]` ...
types: A list of `Tensor` of `int32`, with shapes
`[num_nodes * count1]`, `[num_nodes * count1 * count2]` ...
"""
edge_types = [type_ops.get_edge_type_id(edge_type)
for edge_type in edge_types]
neighbors_list = [tf.reshape(nodes, [-1])]
weights_list = []
type_list = []
neighbors, weights, types = _sample_fanout(
neighbors_list[-1],
edge_types, counts,
default_node=default_node,
N=len(counts))
neighbors_list.extend([tf.reshape(n, [-1]) for n in neighbors])
weights_list.extend([tf.reshape(w, [-1]) for w in weights])
type_list.extend([tf.reshape(t, [-1]) for t in types])
return neighbors_list, weights_list, type_list
def sample_neighbor_layerwise(nodes, edge_types, count,
default_node=-1, weight_func=''):
edge_types = type_ops.get_edge_type_id(edge_types)
res = _sample_neighbor_layerwise_with_adj(nodes, edge_types, count,
weight_func, default_node)
return res[0], tf.SparseTensor(*res[1:4])
def get_top_k_neighbor(nodes, edge_types, k, default_node=-1, condition=''):
edge_types = type_ops.get_edge_type_id(edge_types)
return _get_top_k_neighbor(nodes, edge_types, k, default_node, condition)
def sample_neighbor(nodes, edge_types, count, default_node=-1, condition=''):
edge_types = type_ops.get_edge_type_id(edge_types)
return _sample_neighbor(nodes, edge_types, count, default_node, condition)
def sparse_get_adj(nodes, nb_nodes, edge_types, n=-1, m=-1):
edge_types = type_ops.get_edge_type_id(edge_types)
res = base._LIB_OP.sparse_get_adj(nodes, nb_nodes, edge_types, n, m)
return tf.SparseTensor(*res[:3])