def sgc(x, edge_index, edge_weight, K, kernel, bias=None, renorm=True, improved=False, cache=None):
"""
Functional API for Simple Graph Convolution (SGC).
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param K: Number of hops.(default: :obj:`1`)
:param kernel: Tensor, shape: [num_features, num_output_features], weight.
:param bias: Tensor, shape: [num_output_features], bias.
:param renorm: Whether use renormalization trick (https://arxiv.org/pdf/1609.02907.pdf).
:param improved: Whether use improved GCN or not.
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
:return: Updated node features (x), shape: [num_nodes, num_features]
"""
updated_edge_index, normed_edge_weight = gcn_norm_edge(edge_index, x.shape[0], edge_weight,
renorm, improved, cache)
h = x
for _ in range(K):
h = aggregate_neighbors(
h,
updated_edge_index,
normed_edge_weight,
gcn_mapper,
sum_reducer,
identity_updater
)
h = h @ kernel
if bias is not None:
h += bias
return h
def gcn(x, edge_index, edge_weight, kernel, bias=None, activation=None, improved=False, cache=None):
"""
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param kernel: Tensor, shape: [num_features, num_output_features], weight
:param bias: Tensor, shape: [num_output_features], bias
:param activation: Activation function to use.
:param improved: Whether use improved GCN or not.
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
updated_edge_index, normed_edge_weight = gcn_norm_edge(edge_index, x.shape[0], edge_weight,
improved, cache)
x = x @ kernel
h = aggregate_neighbors(
x, updated_edge_index, normed_edge_weight,
gcn_mapper,
sum_reducer,
identity_updater
)
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
def gin(x,
edge_index,
edge_weight,
mlp_model,
eps=0.0,
activation=None,
cache=None):
"""
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param eps: float, optional, (default: :obj:`0.`).
:param mlp_model: A neural network (multi-layer perceptrons).
:param kernel: Tensor, shape: [num_features, num_output_features], weight
:param cache: A dict for caching A' for GIN. Different graph should not share the same cache dict.
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
h = aggregate_neighbors(x, edge_index, edge_weight, identity_mapper,
sum_reducer, identity_updater)
h = gin_updater(x, h, eps)
h = mlp_model(h)
if activation is not None:
h = activation(h)
return h
def chebynet(x,
edge_index,
edge_weight,
k,
kernels,
bias=None,
activation=None,
normalization_type="sym",
use_dynamic_lambda_max=False,
cache=None):
num_nodes = tf.shape(x)[0]
# lambda_max = chebynet_compute_lambda_max(x, edge_index, edge_weight, normalization_type, cache=cache)
norm_edge_index, norm_edge_weight = chebynet_norm_edge(
edge_index,
num_nodes,
edge_weight,
normalization_type,
use_dynamic_lambda_max=use_dynamic_lambda_max,
cache=cache)
T0_x = x
T1_x = x
out = tf.matmul(T0_x, kernels[0])
if k > 1:
T1_x = aggregate_neighbors(x, norm_edge_index, norm_edge_weight,
gcn_mapper, sum_reducer, identity_updater)
out += tf.matmul(T1_x, kernels[1])
for i in range(2, k):
T2_x = aggregate_neighbors(T1_x, norm_edge_index, norm_edge_weight,
gcn_mapper, sum_reducer,
identity_updater) ##L^T_{k-1}(L^)
T2_x = 2.0 * T2_x - T0_x
out += tf.matmul(T2_x, kernels[i])
T0_x, T1_x = T1_x, T2_x
if bias is not None:
out += bias
if activation is not None:
out = activation(out)
return out
def gat(x,
edge_index,
query_kernel,
query_bias,
query_activation,
key_kernel,
key_bias,
key_activation,
kernel,
bias=None,
activation=None,
num_heads=1,
drop_rate=0.0,
training=False):
num_nodes = x.shape[0]
# self-attention
edge_index, edge_weight = add_self_loop_edge(edge_index, num_nodes)
row, col = edge_index
Q = tf.gather(x, row) @ query_kernel + query_bias
Q = query_activation(Q)
K = tf.gather(x, col) @ key_kernel + key_bias
K = key_activation(K)
V = x @ kernel
# xxxxx_ denotes the multi-head style stuff
Q_ = tf.concat(tf.split(Q, num_heads, axis=-1), axis=0)
K_ = tf.concat(tf.split(K, num_heads, axis=-1), axis=0)
V_ = tf.concat(tf.split(V, num_heads, axis=-1), axis=0)
edge_index_ = tf.concat(
[edge_index + i * num_nodes for i in range(num_heads)], axis=1)
att_score_ = tf.reduce_sum(Q_ * K_, axis=-1)
normed_att_score_ = segment_softmax(att_score_, edge_index_[0],
num_nodes * num_heads)
if training and drop_rate > 0.0:
normed_att_score_ = tf.compat.v2.nn.dropout(normed_att_score_,
drop_rate)
h_ = aggregate_neighbors(V_, edge_index_, normed_att_score_, gcn_mapper,
sum_reducer, identity_updater)
h = tf.concat(tf.split(h_, num_heads, axis=0), axis=-1)
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
def gcn(x,
edge_index,
edge_weight,
kernel,
bias=None,
activation=None,
renorm=True,
improved=False,
cache=None):
"""
Functional API for Graph Convolutional Networks.
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param kernel: Tensor, shape: [num_features, num_output_features], weight
:param bias: Tensor, shape: [num_output_features], bias
:param activation: Activation function to use.
:param renorm: Whether use renormalization trick (https://arxiv.org/pdf/1609.02907.pdf).
:param improved: Whether use improved GCN or not.
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
To use @tf_utils.function with gcn, you should cache the noremd edge information before the first call of the gcn.
- (1) If you're using OOP APIs tfg.layers.GCN:
gcn_layer.cache_normed_edge(graph)
- (2) If you're using functional API tfg.nn.gcn:
from tf_geometric.nn.conv.gcn import gcn_cache_normed_edge
gcn_cache_normed_edge(graph)
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
num_nodes = tf.shape(x)[0]
updated_edge_index, normed_edge_weight = gcn_norm_edge(
edge_index, num_nodes, edge_weight, renorm, improved, cache)
x = x @ kernel
h = aggregate_neighbors(x,
updated_edge_index,
normed_edge_weight,
gcn_mapper,
sum_reducer,
identity_updater,
num_nodes=num_nodes)
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
def chebynet(x,
edge_index,
edge_weight,
K,
lambda_max,
kernel,
bias=None,
activation=None,
normalization_type=None):
num_nodes = x.shape[0]
norm_edge_index, norm_edge_weight = chebynet_norm_edge(
edge_index,
num_nodes,
edge_weight,
lambda_max,
normalization_type=normalization_type)
T0_x = x
T1_x = x
out = tf.matmul(T0_x, kernel[0])
if K > 1:
T1_x = aggregate_neighbors(x, norm_edge_index, norm_edge_weight,
gcn_mapper, sum_reducer, identity_updater)
out += tf.matmul(T1_x, kernel[1])
for i in range(2, K):
T2_x = aggregate_neighbors(T1_x, norm_edge_index, norm_edge_weight,
gcn_mapper, sum_reducer,
identity_updater) ##L^T_{k-1}(L^)
T2_x = 2.0 * T2_x - T0_x
out += tf.matmul(T2_x, kernel[i])
T0_x, T1_x = T1_x, T2_x
if bias is not None:
out += bias
if activation is not None:
out = activation(out)
return out
def tagcn(x,
edge_index,
edge_weight,
K,
kernel,
bias=None,
activation=None,
renorm=False,
improved=False,
cache=None):
"""
Functional API for Topology Adaptive Graph Convolutional Network (TAGCN).
:param x: Tensor, shape: [num_nodes, num_features], node features.
:param edge_index: Tensor, shape: [2, num_edges], edge information.
:param edge_weight: Tensor or None, shape: [num_edges].
:param K: Number of hops.(default: :obj:`3`)
:param kernel: Tensor, shape: [num_features, num_output_features], weight.
:param bias: Tensor, shape: [num_output_features], bias.
:param activation: Activation function to use.
:param renorm: Whether use renormalization trick (https://arxiv.org/pdf/1609.02907.pdf).
:param improved: Whether use improved GCN or not.
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
xs = [x]
updated_edge_index, normed_edge_weight = gcn_norm_edge(
edge_index, x.shape[0], edge_weight, renorm, improved, cache)
for k in range(K):
h = aggregate_neighbors(xs[-1], updated_edge_index, normed_edge_weight,
gcn_mapper, sum_reducer, identity_updater)
xs.append(h)
h = tf.concat(xs, axis=-1)
out = h @ kernel
if bias is not None:
out += bias
if activation is not None:
out = activation(out)
return out
def gcn(x, edge_index, edge_weight, kernel, bias=None, activation=None, improved=False, cache=None):
updated_edge_index, normed_edge_weight = gcn_norm_edge(edge_index, x.shape[0], edge_weight,
improved, cache)
x = x @ kernel
h = aggregate_neighbors(
x, updated_edge_index, normed_edge_weight,
gcn_mapper,
sum_reducer,
identity_updater
)
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
def gin(x, edge_index, edge_weight, mlp_model, eps=0.0, training=None):
"""
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param mlp_model: A neural network (multi-layer perceptrons).
:param eps: float, optional, (default: :obj:`0.`).
:param training: Whether currently executing in training or inference mode.
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
h = aggregate_neighbors(x, edge_index, edge_weight, identity_mapper,
sum_reducer, identity_updater)
h = gin_updater(x, h, eps)
h = mlp_model(h, training=training)
return h
def appnp(x,
edge_index,
edge_weight,
kernels,
biases,
dense_activation=tf.nn.relu,
activation=None,
num_iterations=2,
alpha=0.15,
dense_drop_rate=0.0,
edge_drop_rate=0.0,
cache=None,
training=False):
"""
Functional API for Approximate Personalized Propagation of Neural Predictions (APPNP).
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param kernels: List[Tensor], shape of each Tensor: [num_features, num_output_features], weights
:param biases: List[Tensor], shape of each Tensor: [num_output_features], biases
:param dense_activation: Activation function to use for the dense layers,
except for the last dense layer, which will not be activated.
:param activation: Activation function to use for the output.
:param num_iterations: Number of propagation power iterations.
:param alpha: Teleport Probability.
:param dense_drop_rate: Dropout rate for the input of every dense layer.
:param edge_drop_rate: Dropout rate for the edges/adj used for propagation.
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
To use @tf_utils.function with gcn, you should cache the noremd edge information before the first call of the gcn.
- (1) If you're using OOP APIs tfg.layers.GCN:
gcn_layer.cache_normed_edge(graph)
- (2) If you're using functional API tfg.nn.gcn:
from tf_geometric.nn.conv.gcn import gcn_cache_normed_edge
gcn_cache_normed_edge(graph)
:param training: Python boolean indicating whether the layer should behave in
training mode (adding dropout) or in inference mode (doing nothing).
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
num_nodes = tf.shape(x)[0]
updated_edge_index, normed_edge_weight = gcn_norm_edge(edge_index,
num_nodes,
edge_weight,
cache=cache)
num_dense_layers = len(kernels)
h = x
for i, (kernel, bias) in enumerate(zip(kernels, biases)):
if training and dense_drop_rate > 0.0:
h = tf.compat.v2.nn.dropout(h, dense_drop_rate)
h = h @ kernel + bias
if dense_activation is not None and i < num_dense_layers - 1:
h = dense_activation(h)
if training and edge_drop_rate > 0.0:
normed_edge_weight = tf.compat.v2.nn.dropout(normed_edge_weight,
edge_drop_rate)
prop_h = h
for i in range(num_iterations):
prop_h = aggregate_neighbors(prop_h,
updated_edge_index,
normed_edge_weight,
gcn_mapper,
sum_reducer,
identity_updater,
num_nodes=num_nodes)
prop_h = prop_h * (1.0 - alpha) + h * alpha
if activation is not None:
prop_h = activation(prop_h)
return prop_h
def asap(x,
edge_index,
edge_weight,
node_graph_index,
attention_gcn_kernel,
attention_gcn_bias,
attention_query_kernel,
attention_query_bias,
attention_score_kernel,
attention_score_bias,
le_conv_self_kernel,
le_conv_self_bias,
le_conv_aggr_self_kernel,
le_conv_aggr_self_bias,
le_conv_aggr_neighbor_kernel,
le_conv_aggr_neighbor_bias,
K=None,
ratio=None,
le_conv_activation=tf.nn.sigmoid,
drop_rate=0.0,
training=None,
cache=None):
"""
Functional API for ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representation
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param edge_weight: Tensor or None, shape: [num_edges]
:param node_graph_index: Tensor/NDArray, shape: [num_nodes], graph index for each node
:param K: Keep top K targets for each source
:param ratio: Keep num_targets * ratio targets for each source
:param le_conv_activation: Activation to use for node_score before multiplying node_features with node_score
:param training: Python boolean indicating whether the layer should behave in
training mode (adding dropout) or in inference mode (doing nothing).
:param cache: A dict for caching A' for GCN. Different graph should not share the same cache dict.
:return: [pooled_x, pooled_edge_index, pooled_edge_weight, pooled_node_graph_index]
"""
num_nodes = tf.shape(x)[0]
# num_graphs = tf.reduce_max(node_graph_index) + 1
edge_index, edge_weight = remove_self_loop_edge(edge_index, edge_weight)
edge_index_with_self_loop, edge_weight_with_self_loop = add_self_loop_edge(
edge_index, num_nodes=num_nodes, edge_weight=edge_weight)
row_with_self_loop, col_with_self_loop = edge_index_with_self_loop[
0], edge_index_with_self_loop[1]
attention_h = gcn(x,
edge_index,
edge_weight,
attention_gcn_kernel,
attention_gcn_bias,
cache=cache)
# max_pool -> query
attention_query = aggregate_neighbors(attention_h,
edge_index_with_self_loop,
None,
mapper=identity_mapper,
reducer=max_reducer,
updater=identity_updater,
num_nodes=num_nodes)
attention_query = attention_query @ attention_query_kernel + attention_query_bias
repeated_attention_query = tf.gather(attention_query, row_with_self_loop)
repeated_attention_h = tf.gather(attention_h, col_with_self_loop)
attention_score_h = tf.concat(
[repeated_attention_query, repeated_attention_h], axis=-1)
attention_score = attention_score_h @ attention_score_kernel + attention_score_bias
attention_score = tf.nn.leaky_relu(attention_score, alpha=0.2)
normed_attention_score = segment_softmax(attention_score,
row_with_self_loop, num_nodes)
if training and drop_rate > 0:
normed_attention_score = tf.compat.v2.nn.dropout(
normed_attention_score, rate=drop_rate)
# nodes are clusters
cluster_h = aggregate_neighbors(x,
edge_index_with_self_loop,
tf.reshape(normed_attention_score, [-1]),
gcn_mapper,
sum_reducer,
identity_updater,
num_nodes=num_nodes)
node_score = le_conv(cluster_h,
edge_index,
edge_weight,
le_conv_self_kernel,
le_conv_self_bias,
le_conv_aggr_self_kernel,
le_conv_aggr_self_bias,
le_conv_aggr_neighbor_kernel,
le_conv_aggr_neighbor_bias,
activation=None)
topk_node_index = topk_pool(node_graph_index, node_score, K=K, ratio=ratio)
topk_node_score = tf.gather(node_score, topk_node_index)
if le_conv_activation is not None:
topk_node_score = le_conv_activation(topk_node_score)
pooled_x = tf.gather(cluster_h, topk_node_index) * topk_node_score
num_clusters = tf.shape(topk_node_index)[0]
# node->cluster
cluster_reverse_index = tf.cast(tf.fill([num_nodes], -1), tf.int32)
cluster_reverse_index = tf.tensor_scatter_nd_update(
cluster_reverse_index, tf.expand_dims(topk_node_index, axis=-1),
tf.range(num_clusters))
# row, col = edge_index[0], edge_index[1]
assign_row = tf.gather(cluster_reverse_index, row_with_self_loop)
assign_mask = tf.greater_equal(assign_row, 0)
assign_row = tf.boolean_mask(assign_row, assign_mask)
assign_col = tf.boolean_mask(col_with_self_loop, assign_mask)
assign_edge_index = tf.stack([assign_row, assign_col], axis=0)
assign_edge_weight = tf.boolean_mask(normed_attention_score, assign_mask)
assign_edge_weight = tf.reshape(assign_edge_weight, [-1])
assign_edge_weight = tf.stop_gradient(assign_edge_weight)
# Coarsen in a large BatchGraph.
_, pooled_edge_index, pooled_edge_weight = cluster_pool(
None,
edge_index_with_self_loop,
edge_weight_with_self_loop,
assign_edge_index,
assign_edge_weight,
num_clusters,
num_nodes=num_nodes)
pooled_edge_index, pooled_edge_weight = remove_self_loop_edge(
pooled_edge_index, pooled_edge_weight)
pooled_edge_index, pooled_edge_weight = add_self_loop_edge(
pooled_edge_index, num_clusters, pooled_edge_weight)
pooled_node_graph_index = tf.gather(node_graph_index, topk_node_index)
return pooled_x, pooled_edge_index, pooled_edge_weight, pooled_node_graph_index
def gat(x,
edge_index,
query_kernel,
query_bias,
query_activation,
key_kernel,
key_bias,
key_activation,
kernel,
bias=None,
activation=None,
num_heads=1,
drop_rate=0.0,
training=False):
"""
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param query_kernel: Tensor, shape: [num_features, num_query_features], weight for Q in attention
:param query_bias: Tensor, shape: [num_query_features], bias for Q in attention
:param query_activation: Activation function for Q in attention.
:param key_kernel: Tensor, shape: [num_features, num_key_features], weight for K in attention
:param key_bias: Tensor, shape: [num_key_features], bias for K in attention
:param key_activation: Activation function for K in attention.
:param kernel: Tensor, shape: [num_features, num_output_features], weight
:param bias: Tensor, shape: [num_output_features], bias
:param activation: Activation function to use.
:param num_heads: Number of attention heads.
:param drop_rate: Dropout rate.
:param training: Python boolean indicating whether the layer should behave in
training mode (adding dropout) or in inference mode (doing nothing).
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
num_nodes = x.shape[0]
# self-attention
edge_index, edge_weight = add_self_loop_edge(edge_index, num_nodes)
row, col = edge_index
Q = query_activation(x @ query_kernel + query_bias)
Q = tf.gather(Q, row)
K = key_activation(x @ key_kernel + key_bias)
K = tf.gather(K, col)
V = x @ kernel
# xxxxx_ denotes the multi-head style stuff
Q_ = tf.concat(tf.split(Q, num_heads, axis=-1), axis=0)
K_ = tf.concat(tf.split(K, num_heads, axis=-1), axis=0)
V_ = tf.concat(tf.split(V, num_heads, axis=-1), axis=0)
edge_index_ = tf.concat(
[edge_index + i * num_nodes for i in range(num_heads)], axis=1)
att_score_ = tf.reduce_sum(Q_ * K_, axis=-1)
normed_att_score_ = segment_softmax(att_score_, edge_index_[0],
num_nodes * num_heads)
if training and drop_rate > 0.0:
normed_att_score_ = tf.compat.v2.nn.dropout(normed_att_score_,
drop_rate)
h_ = aggregate_neighbors(V_, edge_index_, normed_att_score_, gcn_mapper,
sum_reducer, identity_updater)
h = tf.concat(tf.split(h_, num_heads, axis=0), axis=-1)
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
def gat(x,
edge_index,
query_kernel,
query_bias,
query_activation,
key_kernel,
key_bias,
key_activation,
kernel,
bias=None,
activation=None,
num_heads=1,
split_value_heads=True,
drop_rate=0.0,
training=False):
"""
:param x: Tensor, shape: [num_nodes, num_features], node features
:param edge_index: Tensor, shape: [2, num_edges], edge information
:param query_kernel: Tensor, shape: [num_features, num_query_features], weight for Q in attention
:param query_bias: Tensor, shape: [num_query_features], bias for Q in attention
:param query_activation: Activation function for Q in attention.
:param key_kernel: Tensor, shape: [num_features, num_key_features], weight for K in attention
:param key_bias: Tensor, shape: [num_key_features], bias for K in attention
:param key_activation: Activation function for K in attention.
:param kernel: Tensor, shape: [num_features, num_output_features], weight
:param bias: Tensor, shape: [num_output_features], bias
:param activation: Activation function to use.
:param num_heads: Number of attention heads.
:param split_value_heads: Boolean. If true, split V as value attention heads, and then concatenate them as output.
Else, num_heads replicas of V are used as value attention heads, and the mean of them are used as output.
:param drop_rate: Dropout rate.
:param training: Python boolean indicating whether the layer should behave in
training mode (adding dropout) or in inference mode (doing nothing).
:return: Updated node features (x), shape: [num_nodes, num_output_features]
"""
num_nodes = tf.shape(x)[0]
# self-attention
edge_index, edge_weight = add_self_loop_edge(edge_index, num_nodes)
row, col = edge_index[0], edge_index[1]
Q = x @ query_kernel + query_bias
if query_activation is not None:
Q = query_activation(Q)
Q = tf.gather(Q, row)
K = x @ key_kernel + key_bias
if key_activation is not None:
K = key_activation(K)
K = tf.gather(K, col)
V = x @ kernel
# xxxxx_ denotes the multi-head style stuff
Q_ = tf.concat(tf.split(Q, num_heads, axis=-1), axis=0)
K_ = tf.concat(tf.split(K, num_heads, axis=-1), axis=0)
# splited queries and keys are modeled as virtual vertices
qk_edge_index_ = tf.concat(
[edge_index + i * num_nodes for i in range(num_heads)], axis=1)
scale = tf.math.sqrt(tf.cast(tf.shape(Q_)[-1], tf.float32))
att_score_ = tf.reduce_sum(Q_ * K_ / scale, axis=-1)
normed_att_score_ = segment_softmax(att_score_, qk_edge_index_[0],
num_nodes * num_heads)
if training and drop_rate > 0.0:
normed_att_score_ = tf.compat.v2.nn.dropout(normed_att_score_,
drop_rate)
if split_value_heads:
V_ = tf.concat(tf.split(V, num_heads, axis=-1), axis=0)
edge_index_ = qk_edge_index_
else:
V_ = V
edge_index_ = tf.tile(edge_index, [1, num_heads])
h_ = aggregate_neighbors(V_, edge_index_, normed_att_score_, gcn_mapper,
sum_reducer, identity_updater)
if split_value_heads:
h = tf.concat(tf.split(h_, num_heads, axis=0), axis=-1)
else:
h = h_ / num_heads
if bias is not None:
h += bias
if activation is not None:
h = activation(h)
return h
