diffalign.models.encoder.egnnΒΆ

Functions

unsorted_segment_sum(data, segment_ids, ...)

Custom PyTorch op to replicate TensorFlow's unsorted_segment_sum.

Classes

EGNN(in_node_nf, in_edge_nf, hidden_nf[, ...])

EGNN_old(in_node_nf, in_edge_nf, hidden_nf)

E_GCL(input_nf, output_nf, hidden_nf[, ...])

Graph Neural Net with global state and fixed number of nodes per graph.

GNN(in_node_nf, in_edge_nf, hidden_nf[, ...])

MLP(dim_in, dim_out, width, nb_layers[, ...])

TransformerNN(in_node_nf, in_edge_nf, hidden_nf)