encoder

encoder#

Encoder module for PrxteinMPNN.

This module contains the EncoderLayer and Encoder classes.

class prxteinmpnn.model.encoder.EncoderLayer(node_features, edge_features, hidden_features, *, key)[source]#

Bases: Module

A single encoder layer for the ProteinMPNN model.

Parameters:

  • node_features (int)

  • edge_features (int)

  • hidden_features (int)

  • key (PRNGKeyArray)

node_features_dim: int#
edge_features_dim: int#
edge_message_mlp: MLP#
norm1: LayerNorm#
dense: MLP#
norm2: LayerNorm#
edge_update_mlp: MLP#
norm3: LayerNorm#
_get_mlp_input(h, e, neighbor_indices)[source]#

Return the input tensor [h_i, e_ij, h_j] for edge_message_mlp.

Return type:

Array

Parameters:

  • h (Int[Array, 'num_atoms num_features'])

  • e (Float[Array, 'num_atoms num_neighbors num_features'])

  • neighbor_indices (Int[Array, 'num_atoms num_neighbors'])

class prxteinmpnn.model.encoder.Encoder(node_features, edge_features, hidden_features, num_layers=3, *, key)[source]#

Bases: Module

The complete encoder module for ProteinMPNN.

Parameters:

  • node_features (int)

  • edge_features (int)

  • hidden_features (int)

  • num_layers (int)

  • key (PRNGKeyArray)

node_feature_dim: int#
layers: tuple[EncoderLayer, ...]#
Contents