lloca.equivectors.lgatr

Edge convolution with L-GATr.

Functions

get_qk_product(q_mv, k_mv, q_s, k_s, edge_index)

Classes

LGATrVectors(n_vectors, num_scalars, ...[, ...])

Initialize internal Module state, shared by both nn.Module and ScriptModule.

class lloca.equivectors.lgatr.LGATrVectors(n_vectors, num_scalars, hidden_mv_channels, hidden_s_channels, net, operation='add', nonlinearity='softmax', aggr='sum', layer_norm=False, lgatr_norm=True, use_amp=False, attention_backend='xformers')

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(fourmomenta, scalars=None, ptr=None, **kwargs)

Runs the forward pass of the module.

message(edge_index, fm_i, fm_j, node_ptr, batch, prefactor)

Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index. This function can take any argument as input which was initially passed to propagate(). Furthermore, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.