core.models.escaip.modules.readout_block

Classes

ReadoutBlock

Readout from each graph attention block for energy and force output

Module Contents

class core.models.escaip.modules.readout_block.ReadoutBlock(global_cfg: fairchem.core.models.escaip.configs.GlobalConfigs, gnn_cfg: fairchem.core.models.escaip.configs.GraphNeuralNetworksConfigs, reg_cfg: fairchem.core.models.escaip.configs.RegularizationConfigs)

Bases: torch.nn.Module

Readout from each graph attention block for energy and force output

backbone_dtype

energy_reduce

use_edge_readout

use_global_readout

node_ffn

pre_node_norm

forward(data, node_features, edge_features)

Output:

Global Readout (G, H); Node Readout (N, H); Edge Readout (N, max_nei, H)