protein_workshop.models#

To switch between different encoder architectures, simply change the encoder argument in the launch command. For example:

workshop train encoder=<encoder_name> dataset=cath task=inverse_folding trainer=cpu
# or
python proteinworkshop/train.py encoder=<encoder_name> dataset=cath task=inverse_folding trainer=cpu # or trainer=gpu

Where <encoder_name> is given by bracketed name in the listing below. For example, the encoder name for SchNet is schnet.

Note

To change encoder hyperparameters, either

Edit the config file directly, or
Provide commands in the form:

workshop train encoder=<encoder_name> encoder.num_layer=3 encoder.readout=mean dataset=cath task=inverse_folding trainer=cpu
# or
python proteinworkshop/train.py encoder=<encoder_name> encoder.num_layer=3 encoder.readout=mean dataset=cath task=inverse_folding trainer=cpu # or trainer=gpu

The following structural encoders are currently supported:

Base Classes #

class proteinworkshop.models.base.BaseModel(*args: Any, **kwargs: Any)[source]#

Bases: LightningModule, ABC

compute_loss(y_hat: ModelOutput, y: Label) → Dict[str, Tensor][source]#

Compute loss by iterating over all outputs.

In the case of multiple losses, the total loss is also included in the output dictionary of losses.

Parameters:

y_hat (ModelOutput) – Output of model. This should be a dictionary of outputs (torch.Tensor) indexed by the output name (str)
y (Label) – Labels. This should be a dictionary of labels (torch.Tensor) indexed by the output name (str)

Returns:

Dictionary of losses indexed by output name (str)

Return type:

Dict[str, torch.Tensor]

config: DictConfig#

configure_losses(loss_dict: Dict[str, str]) → Dict[str, Callable][source]#

Configures losses from config. Returns a dictionary of losses mapping each output name to its respective loss function.

Parameters:: loss_dict (Dict[str, str]) – Config dictionary of losses indexed by output name
Returns:: Dictionary of losses indexed by output name
Return type:: Dict[str, Callable]

configure_metrics()[source]#

Instantiates metrics from config.

Metrics are Torchmetrics Objects torchmetrics.Metric (see torchmetrics)

Metrics are set as model attributes as:

{stage}_{output}_{metric_name} (e.g. train_residue_type_f1_score)

configure_optimizers()[source]#

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Returns:

Any of these 6 options.

Single optimizer.
List or Tuple of optimizers.
Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).
Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.
None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

# The ReduceLROnPlateau scheduler requires a monitor
def configure_optimizers(self):
    optimizer = Adam(...)
    return {
        "optimizer": optimizer,
        "lr_scheduler": {
            "scheduler": ReduceLROnPlateau(optimizer, ...),
            "monitor": "metric_to_track",
            "frequency": "indicates how often the metric is updated"
            # If "monitor" references validation metrics, then "frequency" should be set to a
            # multiple of "trainer.check_val_every_n_epoch".
        },
    }


# In the case of two optimizers, only one using the ReduceLROnPlateau scheduler
def configure_optimizers(self):
    optimizer1 = Adam(...)
    optimizer2 = SGD(...)
    scheduler1 = ReduceLROnPlateau(optimizer1, ...)
    scheduler2 = LambdaLR(optimizer2, ...)
    return (
        {
            "optimizer": optimizer1,
            "lr_scheduler": {
                "scheduler": scheduler1,
                "monitor": "metric_to_track",
            },
        },
        {"optimizer": optimizer2, "lr_scheduler": scheduler2},
    )

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

Lightning calls .backward() and .step() automatically in case of automatic optimization.
If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.
If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.
If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.
If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.
If you need to control how often the optimizer steps, override the optimizer_step() hook.

featurise(batch: Batch | ProteinBatch) → Batch | ProteinBatch[source]#

Applies the featuriser (self.featuriser) to a batch of data.

Invariant Encoders #

SchNet (`schnet`)#

class proteinworkshop.models.graph_encoders.schnet.SchNetModel(hidden_channels: int = 128, out_dim: int = 1, num_filters: int = 128, num_layers: int = 6, num_gaussians: int = 50, cutoff: float = 10, max_num_neighbors: int = 32, readout: str = 'add', dipole: bool = False, mean: float | None = None, std: float | None = None, atomref: Tensor | None = None)[source]#

Bases: SchNet

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the SchNet encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x: Node features (shape: \((n, d)\))
pos: Node positions (shape: \((n, 3)\))
edge_index: Edge indices (shape: \((2, e)\))
batch: Batch indices (shape: \((n,)\))

Returns:: Set of required batch attributes
Return type:: Set[str]

training: bool#

DimeNet++ (`dimenet_plus_plus`)#

class proteinworkshop.models.graph_encoders.dimenetpp.DimeNetPPModel(hidden_channels: int = 128, out_dim: int = 1, num_layers: int = 4, int_emb_size: int = 64, basis_emb_size: int = 8, out_emb_channels: int = 256, num_spherical: int = 7, num_radial: int = 6, cutoff: float = 10, max_num_neighbors: int = 32, envelope_exponent: int = 5, num_before_skip: int = 1, num_after_skip: int = 2, num_output_layers: int = 3, act: str | Callable = 'swish', readout: str = 'add')[source]#

Bases: DimeNetPlusPlus

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the DimeNet++ encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x: Node features (shape: \((n, d)\))
pos: Node positions (shape: \((n, 3)\))
edge_index: Edge indices (shape: \((2, e)\))
batch: Batch indices (shape: \((n,)\))

Returns:: _description_
Return type:: Set[str]

training: bool#

GearNet (`gear_net`, `gear_net_edge`)#

class proteinworkshop.models.graph_encoders.gear_net.GearNet(input_dim: int, num_relation: int, num_layers: int, emb_dim: int, short_cut: bool, concat_hidden: bool, batch_norm: bool, num_angle_bin: int | None, activation: str = 'relu', pool: str = 'sum')[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the GearNet encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

gear_net_edge_features(b: Batch | ProteinBatch) → Tensor[source]#

Compute edge features for the gear net encoder.

Concatenate node features of the two nodes in each edge
Concatenate the edge type
Compute the distance between the two nodes in each edge
Compute the sequence distance between the two nodes in each edge

Parameters:: b (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Edge features
Return type:: torch.Tensor

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x Positions (shape [num_nodes, 3])
edge_index Edge indices (shape [2, num_edges])
edge_type Edge types (shape [num_edges])
edge_attr Edge attributes (shape [num_edges, num_edge_features])
num_nodes Number of nodes (int)
batch Batch indices (shape [num_nodes])

Returns:: Set of required batch attributes.
Return type:: Set[str]

training: bool#

CDConv (`cdconv`)#

CDConv implementation adapted from the MIT-licensed source: https://github.com/hehefan/Continuous-Discrete-Convolution/tree/main

class proteinworkshop.models.graph_encoders.cdconv.CDConvModel(geometric_radius: float, sequential_kernel_size: float, kernel_channels: List[int], channels: List[int], base_width: float = 16.0, embedding_dim: int = 16, batch_norm: bool = True, dropout: float = 0.2, bias: bool = False, readout: str = 'sum')[source]#

Bases: PatchedModule

forward(data: ProteinBatch | Batch) → EncoderOutput[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

static orientation(pos)[source]#

property required_batch_attributes: Set[str]#

training: bool#

Vector-Equivariant Encoders #

EGNN (`egnn`)#

class proteinworkshop.models.graph_encoders.egnn.EGNNModel(num_layers: int = 5, emb_dim: int = 128, activation: str = 'relu', norm: str = 'layer', aggr: str = 'mean', pool: str = 'mean', residual: bool = True, dropout: float = 0.1)[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the EGNN encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x: Node features (shape: \((n, d)\))
pos: Node positions (shape: \((n, 3)\))
edge_index: Edge indices (shape: \((2, e)\))
batch: Batch indices (shape: \((n,)\))

Returns:: Set of required batch attributes
Return type:: Set[str]

training: bool#

GVP (`gvp`)#

class proteinworkshop.models.graph_encoders.gvp.GVPGNNModel(s_dim: int = 128, v_dim: int = 16, s_dim_edge: int = 32, v_dim_edge: int = 1, r_max: float = 10.0, num_bessel: int = 8, num_polynomial_cutoff: int = 5, num_layers: int = 5, pool: str = 'sum', residual: bool = True)[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the GVP-GNN encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

edge_index (shape [2, num_edges])
pos (shape [num_nodes, 3])
x (shape [num_nodes, num_node_features])
batch (shape [num_nodes])

Returns:: _description_
Return type:: Set[str]

training: bool#

GCPNet (`gcpnet`)#

class proteinworkshop.models.graph_encoders.gcpnet.GCPNetModel(num_layers: int = 5, node_s_emb_dim: int = 128, node_v_emb_dim: int = 16, edge_s_emb_dim: int = 32, edge_v_emb_dim: int = 4, r_max: float = 10.0, num_rbf: int = 8, activation: str = 'silu', pool: str = 'sum', **kwargs)[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the GCPNet encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: List[str]#

training: bool#

Tensor-Equivariant Encoders #

TFN (`tfn`)#

class proteinworkshop.models.graph_encoders.tfn.TensorProductModel(r_max: float = 10.0, num_basis: int = 16, max_ell: int = 2, num_layers: int = 4, hidden_irreps='32x0e + 32x0o + 8x1e + 8x1o + 4x2e + 4x2o', mlp_dim: int = 256, aggr: str = 'mean', pool: str = 'sum', residual: bool = True, batch_norm: bool = True, gate: bool = False, dropout: float = 0.1)[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the TFN encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x: Node features (shape: \((n, d)\))
pos: Node positions (shape: \((n, 3)\))
edge_index: Edge indices (shape: \((2, e)\))
batch: Batch indices (shape: \((n,)\))

Returns:: Set of required batch attributes
Return type:: Set[str]

training: bool#

Multi-Atomic Cluster Expansion (`mace`)#

class proteinworkshop.models.graph_encoders.mace.MACEModel(r_max: float = 10.0, num_bessel: int = 8, num_polynomial_cutoff: int = 5, max_ell: int = 2, correlation: int = 3, num_layers: int = 2, hidden_irreps='32x0e + 32x1o + 32x2e', mlp_dim: int = 256, aggr: str = 'mean', pool: str = 'sum', residual: bool = True, batch_norm: bool = True, gate: bool = False, dropout: float = 0.1)[source]#

Bases: PatchedModule

forward(batch: Batch | ProteinBatch) → EncoderOutput[source]#

Implements the forward pass of the MACE encoder.

Returns the node embedding and graph embedding in a dictionary.

Parameters:: batch (Union[Batch, ProteinBatch]) – Batch of data to encode.
Returns:: Dictionary of node and graph embeddings. Contains node_embedding and graph_embedding fields. The node embedding is of shape \((|V|, d)\) and the graph embedding is of shape \((n, d)\), where \(|V|\) is the number of nodes and \(n\) is the number of graphs in the batch and \(d\) is the dimension of the embeddings.
Return type:: EncoderOutput

property required_batch_attributes: Set[str]#

Required batch attributes for this encoder.

x: Node features (shape: \((n, d)\))
pos: Node positions (shape: \((n, 3)\))
edge_index: Edge indices (shape: \((2, e)\))
batch: Batch indices (shape: \((n,)\))

Returns:: Set of required batch attributes
Return type:: Set[str]

training: bool#

Sequence-Based Encoders #

Evolutionary Scale Modeling (`esm`)#

Modified from TorchDrug.

class proteinworkshop.models.graph_encoders.esm_embeddings.EvolutionaryScaleModeling(path: str | PathLike, model: str = 'ESM-2-650M', readout: str = 'mean', mlp_post_embed: bool = True, dropout: float = 0.1, finetune: bool = False)[source]#

Bases: PatchedModule

The protein language model, Evolutionary Scale Modeling (ESM) proposed in Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences.

Parameters:

(str) (readout) – path to store ESM model weights
(str) – model name. Available model names are ESM-1b, ESM-1v and ESM-1b-regression.
(str) – readout function. Available functions are sum and mean.
(bool) (finetune) – whether to use MLP to combine ESM embeddings with input features
(float) (dropout) – dropout rate for MLP
(bool) – whether to finetune ESM model

esm_embed(batch: Batch | ProteinBatch, device: device | str | None = None) → Tensor[source]#: Compute residue ESM embeddings for input proteins

forward(batch: Batch | ProteinBatch, device: device | str | None = None) → EncoderOutput[source]#

Compute the residue representations and the graph representation(s).

Parameters:

(Protein) (graph) – \(n\) protein(s)
(Tensor) (input) – input node representations
optional) (device (torch.device or str,) – device on which to compute and update representations

Returns:

dict with residue_feature and graph_feature fields: residue representations of shape \((|V_{res}|, d)\), graph representations of shape \((n, d)\)

load_weight(path: str, model: str) → Tuple[PatchedModule, Alphabet][source]#

Load ESM model weights and their corresponding alphabet.

Parameters:

(str) (model) – path to store ESM model weights
(str) – model name. Available model names are ESM-1b, ESM-1v and ESM-1b-regression.

Returns:

ESM model and its alphabet as nn.Module and esm.data.Alphabet objects, respectively.

max_input_length = 1022#

md5: Dict[str, str] = {'ESM-1b': 'ba8914bc3358cae2254ebc8874ee67f6', 'ESM-1b-regression': 'e7fe626dfd516fb6824bd1d30192bdb1', 'ESM-1v': '1f04c2d2636b02b544ecb5fbbef8fefd', 'ESM-2-150M': '229fcf8f9f3d4d442215662ca001b906', 'ESM-2-15B': 'af61a9c0b792ae50e244cde443b7f4ac', 'ESM-2-35M': 'a894ddb31522e511e1273abb23b5f974', 'ESM-2-3B': 'd37a0d0dbe7431e48a72072b9180b16b', 'ESM-2-650M': 'ba6d997e29db07a2ad9dca20e024b102', 'ESM-2-8M': '8039fc9cee7f71cd2633b13b5a38ff50'}#

model_names: List[str] = ['ESM-1b', 'ESM-1v', 'ESM-1b-regression', 'ESM-2-8M', 'ESM-2-35M', 'ESM-2-150M', 'ESM-2-650M', 'ESM-2-3B', 'ESM-2-15B']#

num_layer: Dict[str, int] = {'ESM-1b': 33, 'ESM-1v': 33, 'ESM-2-150M': 30, 'ESM-2-15B': 48, 'ESM-2-35M': 12, 'ESM-2-3B': 36, 'ESM-2-650M': 33, 'ESM-2-8M': 6}#

output_dim: Dict[str, int] = {'ESM-1b': 1280, 'ESM-1v': 1280, 'ESM-2-150M': 640, 'ESM-2-15B': 5120, 'ESM-2-35M': 480, 'ESM-2-3B': 2560, 'ESM-2-650M': 1280, 'ESM-2-8M': 320}#

property required_batch_attributes: Set[str]#

Return the requied attributes for each batch.

Returns:: set of required attributes

url: Dict[str, str] = {'ESM-1b': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm1b_t33_650M_UR50S.pt', 'ESM-1b-regression': 'https://dl.fbaipublicfiles.com/fair-esm/regression/esm1b_t33_650M_UR50S-contact-regression.pt', 'ESM-1v': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_1.pt', 'ESM-2-150M': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt', 'ESM-2-15B': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t48_15B_UR50D.pt', 'ESM-2-35M': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt', 'ESM-2-3B': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt', 'ESM-2-650M': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt', 'ESM-2-8M': 'https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt'}#

Decoders #

Linear Decoders

class proteinworkshop.models.decoders.mlp_decoder.LinearSkipBlock(hidden_dim: List[int], activations: List[Literal['relu', 'elu', 'leaky_relu', 'tanh', 'sigmoid', 'none', 'silu', 'swish']], out_dim: int, dropout: float = 0.0, skip: Literal['sum', 'concat'] = 'sum')[source]#

Bases: PatchedModule

forward(x: Tensor) → Tensor[source]#

Implements the forward pass of the MLP decoder with skip connections.

Parameters:: x (torch.Tensor) – Input tensor
Returns:: Output tensor
Return type:: torch.Tensor

training: bool#

class proteinworkshop.models.decoders.mlp_decoder.MLPDecoder(hidden_dim: List[int], out_dim: int, activations: List[Literal['relu', 'elu', 'leaky_relu', 'tanh', 'sigmoid', 'none', 'silu', 'swish']], dropout: float, skip: bool = True, input: str | None = None)[source]#

Bases: PatchedModule

forward(x: Tensor) → Tensor[source]#

Forward pass of MLP decoder.

Parameters:: x (torch.Tensor) – Input tensor
Returns:: Output tensor
Return type:: torch.Tensor

training: bool#

class proteinworkshop.models.decoders.mlp_decoder.PositionDecoder(num_message_layers: int, message_hidden_dim: int, message_activation: Literal['relu', 'elu', 'leaky_relu', 'tanh', 'sigmoid', 'none', 'silu', 'swish'], message_dropout: float, message_skip: bool, num_distance_layers: int, distance_hidden_dim: int, distance_activation: Literal['relu', 'elu', 'leaky_relu', 'tanh', 'sigmoid', 'none', 'silu', 'swish'], distance_dropout: float, distance_skip: bool, aggr: str = 'sum')[source]#

Bases: PatchedModule

forward(edge_index: Tensor, scalar_features: Tensor, pos: Tensor) → Tensor[source]#

Implement forward pass of MLP decoder for equivariant position prediction.

Parameters:

edge_index (torch.Tensor) – Tensor of edge indices
scalar_features (torch.Tensor) – Tensor of scalar features (N x D)
pos (torch.Tensor) – Tensor of positions (N x 3)

Returns:

Tensor of predicted positions (N x 3)

Return type:

torch.Tensor

training: bool#