// graph_attention.mbt -- Graph Attention Network (v0.139.0).
//
// GAT (Velickovic et al. 2018) replaces the fixed mean/sum aggregation
// of GCN with a learned attention coefficient per edge:
//
// e_uv = a_L^T [W h_u || W h_v]
// alpha_uv = softmax_v(e_uv) (softmax over the IN-neighbours)
// h'_v = sigma( sum_u alpha_uv · W h_u )
//
// where `||` is concatenation. The attention head `a` is itself
// learnable, which is what lets GAT focus on important neighbours
// instead of averaging all of them.
//
// Multi-head attention concatenates K independent heads (rather than
// averaging them) to keep the output dimension stable.
//
// Scope of v0.139.0:
// - GATLayer (single head) + gat_layer_forward.
// - GraphAttention (multi-head stack) + graph_attention_forward.
// - LeakyReLU (slope 0.2, the GAT default).
//
// Reference: Velickovic et al. 2018.
///|
/// LeakyReLU with slope 0.2 (the GAT default negative slope).
pub fn gat_leaky_relu(x : Float) -> Float {
if x > 0.0F { x } else { 0.2F * x }
}
///|
/// GATLayer: one attention head mapping in_dim -> out_dim.
pub struct GATLayer {
in_dim : Int
out_dim : Int
// Shared linear projection W: (out_dim x in_dim).
w : Array[Array[Float]]
// Attention vector a, split into a_src and a_dst halves so the
// score e_uv = a_src . Wh_u + a_dst . Wh_v.
a_src : Array[Float]
a_dst : Array[Float]
}
///|
/// Build a GATLayer with xavier-normal init.
pub fn GATLayer::new(
in_dim : Int,
out_dim : Int,
seed : UInt64,
) -> GATLayer {
let std = sqrtf(2.0F / Float::from_int(in_dim))
let rng = Xoshiro::from_state(
seed + 10UL, seed + 11UL, seed + 12UL, seed + 13UL,
)
let w = xavier_normal(out_dim, in_dim, std, rng)
let a_src : Array[Float] = Array::make(out_dim, 0.0F)
let a_dst : Array[Float] = Array::make(out_dim, 0.0F)
for i in 0.. Array[Float] {
let n = g.n_nodes
let d = layer.out_dim
// Project every node: p = W h, flat [n x d].
let p : Array[Float] = Array::make(n * d, 0.0F)
for v in 0..= n || dst < 0 || dst >= n {
logits[e] = 0.0F
} else {
logits[e] = t_src[s] + t_dst[dst]
}
}
let deg = graph_in_degree(g)
let alpha : Array[Float] = Array::make(g.n_edges, 0.0F)
for v in 0.. m {
m = logits[e]
}
}
let mut sum_exp = 0.0F
for e in 0..= n || dst < 0 || dst >= n {
continue
}
let a = alpha[e]
let src_off = s * d
let dst_off = dst * d
for o in 0..= 1
///
/// The per-head qualifier is the whole point. Multi-head attention
/// CONCATENATES its heads, so layer l does not read `dims[l]` -- it
/// reads `num_heads * dims[l]`. Feeding layer l a bare `dims[l]` built
/// a stack whose declared input width was `num_heads` times too small;
/// `gat_layer_forward` then re-strided the concatenated buffer with the
/// layer's own (small) `in_dim` and read the wrong elements, silently,
/// for every layer past the first. With `num_heads == 1` the two agree,
/// which is why this only ever bit multi-layer multi-head stacks --
/// the one shape nothing in the package had built until now.
pub fn GraphAttention::new(
dims : Array[Int],
num_heads : Int,
seed : UInt64,
) -> GraphAttention {
let num_layers = dims.length() - 1
let layers : Array[Array[GATLayer]] = Array::make(num_layers, [])
for l in 0.. Int {
net.layers[l][0].out_dim
}
///|
/// Node-embedding width of a multi-head GAT stack: the heads are
/// concatenated, so the node row is `num_heads` times the last layer's
/// per-head width.
pub fn gat_stack_node_dim(net : GraphAttention) -> Int {
net.num_heads * gat_stack_out_width(net, net.num_layers - 1)
}
///|
/// helper: ELU activation (x > 0 ? x : exp(x) - 1), the GAT hidden
/// activation.
pub fn gat_elu(x : Float) -> Float {
if x > 0.0F { x } else { expf(x) - 1.0F }
}
///|
/// Forward the multi-head GAT stack, caching each layer's input.
/// Returns `(output, inputs)` where `inputs[l]` is the buffer layer l
/// reads -- the concatenated output of layer l-1.
///
/// This is the real implementation; `graph_attention_forward` is a
/// wrapper over it. Having one body rather than two is deliberate: the
/// backward pass slices `d_out` by the same per-layer widths this
/// function writes, and a second copy of the layout arithmetic is
/// exactly how a backward ends up disagreeing with its own forward.
pub fn graph_attention_forward_with_inputs(
net : GraphAttention,
g : Graph,
h : Array[Float],
) -> (Array[Float], Array[Array[Float]]) {
let heads = net.num_heads
let n_layers = net.num_layers
let inputs : Array[Array[Float]] = Array::make(n_layers, [])
let mut acc = h
let mut i = 0
while i < n_layers {
let d = gat_stack_out_width(net, i)
// The LAST layer is linear so its output can be read as logits;
// every earlier layer concatenates its heads and applies ELU.
let last = i == n_layers - 1
inputs[i] = acc
let out : Array[Float] = Array::make(g.n_nodes * heads * d, 0.0F)
for hd in 0.. Array[Float] {
let (out, _) = graph_attention_forward_with_inputs(net, g, h)
out
}