# Copyright Generate Biomedicines, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Layers for building graph neural networks. This module contains layers for building neural networks that can process graph-structured data. The internal representations of these layers are node and edge embeddings. """ from typing import Callable, List, Optional, Tuple import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint from tqdm.autonotebook import tqdm from chroma.layers.attention import Attention class GraphNN(nn.Module): """Graph neural network with optional edge updates. Args: num_layers (int): Number of layers. dim_nodes (int): Hidden dimension of node tensor. dim_edges (int): Hidden dimension of edge tensor. dropout (float): Dropout rate. node_mlp_layers (int): Node update function, number of hidden layers. Default is 1. node_mlp_dim (int): Node update function, hidden dimension. Default is to match MLP output dimension. update_edge (Boolean): Include an edge-update step. Default: True edge_mlp_layers (int): Edge update function, number of hidden layers. Default is 1. edge_mlp_dim (int): Edge update function, hidden dimension. Default is to match MLP output dimension. mlp_activation (str): MLP nonlinearity. `'relu'`: Rectified linear unit. `'softplus'`: Softplus. norm (str): Which normalization function to apply between layers. `'transformer'`: Default layernorm `'layer'`: Masked Layer norm with shape (input.shape[1:]) `'instance'`: Masked Instance norm scale (float): Scaling factor of edge input when updating node (default=1.0) attentional (bool): If True, use attention for message aggregation function instead of a sum. Default is False. num_attention_heads (int): Number of attention heads (if attentional) to use. Default is 4. Inputs: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i (tensor, optional): Node mask with shape `(num_batch, num_nodes)` mask_ij (tensor, optional): Edge mask with shape `(num_batch, num_nodes, num_neighbors)` Outputs: node_h_out (torch.Tensor): Updated node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h_out (torch.Tensor): Updated edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. """ def __init__( self, num_layers: int, dim_nodes: int, dim_edges: int, node_mlp_layers: int = 1, node_mlp_dim: Optional[int] = None, edge_update: bool = True, edge_mlp_layers: int = 1, edge_mlp_dim: Optional[int] = None, mlp_activation: str = "relu", dropout: float = 0.0, norm: str = "transformer", scale: float = 1.0, skip_connect_input: bool = False, attentional: bool = False, num_attention_heads: int = 4, checkpoint_gradients: bool = False, ): super(GraphNN, self).__init__() ## 残差网络 self.skip_connect_input = skip_connect_input """ 优化内存：正常的训练过程中，为了计算梯度，需要存储前向传播中所有层的激活值。 使用梯度检查点时，只在特定层保留这些激活值，并在需要时重新计算它们 """ self.checkpoint_gradients = checkpoint_gradients self.layers = nn.ModuleList( [ GraphLayer( dim_nodes=dim_nodes, dim_edges=dim_edges, node_mlp_layers=node_mlp_layers, node_mlp_dim=node_mlp_dim, edge_update=edge_update, edge_mlp_layers=edge_mlp_layers, edge_mlp_dim=edge_mlp_dim, mlp_activation=mlp_activation, dropout=dropout, norm=norm, scale=scale, attentional=attentional, num_attention_heads=num_attention_heads, ) for _ in range(num_layers) ] ) def forward( self, node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: Optional[torch.Tensor] = None, mask_ij: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # Run every layer sequentially node_h_init = node_h edge_h_init = edge_h for i, layer in enumerate(self.layers): if self.skip_connect_input: node_h = node_h + node_h_init edge_h = edge_h + edge_h_init # Update edge and node node_h, edge_h = self.checkpoint( layer, node_h, edge_h, edge_idx, mask_i, mask_ij ) if self.skip_connect_input: node_h = node_h - node_h_init edge_h = edge_h - edge_h_init # If mask was provided, apply it if mask_i is not None: node_h = node_h * (mask_i.unsqueeze(-1) != 0).type(torch.float32) if mask_ij is not None: edge_h = edge_h * (mask_ij.unsqueeze(-1) != 0).type(torch.float32) return node_h, edge_h def checkpoint(self, layer, *args): if self.checkpoint_gradients: return checkpoint(layer, *args) else: return layer(*args) def sequential( self, tensors: dict, pre_step_function: Callable = None, post_step_function: Callable = None, ) -> dict: """Decode the GNN sequentially along the node index `t`, with callbacks. Args: tensors (dict): Initial set of state tensors. At minimum this should include the arguments to `forward`, namely `node_h`, `edge_h`, `edge_idx`, `mask_i`, and `mask_ij`. pre_step_function (function, optional): Callback function that is optionally applied to `tensors` before each sequential GNN step as `tensors_new = pre_step_function(t, pre_step_function)` where `t` is the node index being updated. It should update elements of the `tensors` dictionary, and it can access and update the intermediate GNN state cache via the keyed lists of tensors in `node_h_cache` and `edge_h_cache`. post_step_function (function, optional): Same as `pre_step_function`, but optionally applied after each sequential GNN step. Returns: tensors (dict): Processed set of tensors. """ # Initialize the state cache tensors["node_h_cache"], tensors["edge_h_cache"] = self.init_steps( tensors["node_h"], tensors["edge_h"] ) # Sequential iteration num_steps = tensors["node_h"].size(1) for t in tqdm(range(num_steps), desc="Sequential decoding"): if pre_step_function is not None: tensors = pre_step_function(t, tensors) tensors["node_h_cache"], tensors["edge_h_cache"] = self.step( t, tensors["node_h_cache"], tensors["edge_h_cache"], tensors["edge_idx"], tensors["mask_i"], tensors["mask_ij"], ) if post_step_function is not None: tensors = post_step_function(t, tensors) return tensors def init_steps( self, node_h: torch.Tensor, edge_h: torch.Tensor ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: """Initialize cached node and edge features. Args: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. Returns: node_h_cache (torch.Tensor): List of cached node features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, dim_nodes)`. edge_h_cache (torch.Tensor): List of cached edge features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. """ num_layers = len(self.layers) node_h_cache = [node_h.clone() for _ in range(num_layers + 1)] edge_h_cache = [edge_h.clone() for _ in range(num_layers + 1)] return node_h_cache, edge_h_cache def step( self, t: int, node_h_cache: List[torch.Tensor], edge_h_cache: List[torch.Tensor], edge_idx: torch.LongTensor, mask_i: Optional[torch.Tensor] = None, mask_ij: Optional[torch.Tensor] = None, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: """Process GNN update for a specific node index t from cached intermediates. Inputs: t (int): Node index to decode. node_h_cache (List[torch.Tensor]): List of cached node features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, dim_nodes)`. edge_h_cache (List[torch.Tensor]): List of cached edge features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i (torch.Tensor, optional): Node mask with shape `(num_batch, num_nodes)`. mask_ij (torch.Tensor, optional): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. Outputs: node_h_cache (List[torch.Tensor]): Updated list of cached node features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, dim_nodes)`. This method updates the tensors in place for memory. edge_h_cache (List[torch.Tensor]): Updated list of cached edge features with `num_layers + 1` tensors of shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. """ if self.skip_connect_input: raise NotImplementedError for i, layer in enumerate(self.layers): # Because the edge updates depend on the updated nodes, # we need both the input node features node_h and also # the previous output node states node_h node_h = node_h_cache[i] node_h_out = node_h_cache[i + 1] edge_h = edge_h_cache[i] # Update edge and node node_h_t, edge_h_t = checkpoint( layer.step, t, node_h, node_h_out, edge_h, edge_idx, mask_i, mask_ij ) # Scatter them in place node_h_cache[i + 1].scatter_( 1, (t * torch.ones_like(node_h_t)).long(), node_h_t ) edge_h_cache[i + 1].scatter_( 1, (t * torch.ones_like(edge_h_t)).long(), edge_h_t ) return node_h_cache, edge_h_cache ## GNNLayer class GraphLayer(nn.Module): """Graph layer that updates each node i given adjacent nodes and edges. Args: dim_nodes (int): Hidden dimension of node tensor. dim_edges (int): Hidden dimension of edge tensor. node_mlp_layers (int): Node update function, number of hidden layers. Default: 1. node_mlp_dim (int): Node update function, hidden dimension. Default: Matches MLP output dimension. update_edge (Boolean): Include an edge-update step. Default: True edge_mlp_layers (int): Edge update function, number of hidden layers. Default: 1. edge_mlp_dim (int): Edge update function, hidden dimension. Default: Matches MLP output dimension. mlp_activation (str): MLP nonlinearity. `'relu'`: Rectified linear unit. `'softplus'`: Softplus. dropout (float): Dropout rate. norm (str): Which normalization function to apply between layers. `'transformer'`: Default layernorm `'layer'`: Masked Layer norm with shape (input.shape[1:]) `'instance'`: Masked Instance norm scale (float): Scaling factor of edge input when updating node (default=1.0) Inputs: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i (tensor, optional): Node mask with shape `(num_batch, num_nodes)` mask_ij (tensor, optional): Edge mask with shape `(num_batch, num_nodes, num_neighbors)` Outputs: node_h_out (torch.Tensor): Updated node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h_out (torch.Tensor): Updated edge features with shape `(num_batch, num_nodes, num_neighbors, dim_nodes)`. """ def __init__( self, dim_nodes: int, dim_edges: int, node_mlp_layers: int = 1, node_mlp_dim: Optional[int] = None, edge_update: bool = True, edge_mlp_layers: int = 1, edge_mlp_dim: Optional[int] = None, mlp_activation: str = "relu", dropout: float = 0.0, norm: str = "transformer", scale: float = 1.0, attentional: bool = False, num_attention_heads: int = 4, ): super(GraphLayer, self).__init__() # Store scale self.scale = scale self.dim_nodes = dim_nodes self.dim_edges = dim_edges self.attentional = attentional self.node_norm_layer = MaskedNorm( dim=1, num_features=dim_nodes, affine=True, norm=norm ) self.message_mlp = MLP( dim_in=2 * dim_nodes + dim_edges, dim_out=dim_nodes, num_layers_hidden=edge_mlp_layers, dim_hidden=edge_mlp_dim, activation=mlp_activation, dropout=dropout, ) self.update_mlp = MLP( dim_in=2 * dim_nodes, dim_out=dim_nodes, num_layers_hidden=node_mlp_layers, dim_hidden=node_mlp_dim, activation=mlp_activation, dropout=dropout, ) self.edge_update = edge_update self.edge_norm_layer = MaskedNorm( dim=2, num_features=dim_edges, affine=True, norm=norm ) if self.edge_update: self.edge_mlp = MLP( dim_in=2 * dim_nodes + dim_edges, dim_out=dim_edges, num_layers_hidden=edge_mlp_layers, dim_hidden=edge_mlp_dim, activation=mlp_activation, dropout=dropout, ) if self.attentional: self.attention = Attention(n_head=num_attention_heads, d_model=dim_nodes) ## attention def attend( self, node_h: torch.Tensor, messages: torch.Tensor, mask_ij: torch.Tensor ) -> torch.Tensor: B, L, K, D = messages.size() queries = node_h.reshape(-1, 1, D) keys = messages.reshape(-1, K, D) values = messages.reshape(-1, K, D) mask = mask_ij.reshape(-1, 1, 1, K).bool() if mask_ij is not None else None return self.attention(queries, keys, values, mask=mask).reshape(B, L, D) ## _normalize:Edge and node def _normalize(self, node_h, edge_h, mask_i=None, mask_ij=None): # Normalize node and edge embeddings node_h_norm = self.node_norm_layer(node_h, mask_i) edge_h_norm = self.edge_norm_layer(edge_h, mask_ij) return node_h_norm, edge_h_norm ## ? def _normalize_t( self, edge_node_stack_t, mask_ij_t, include_nodes=True, include_edges=True ): # Apply normalization (since we have only normalized time t information) node_i_t = edge_node_stack_t[:, :, :, : self.dim_nodes] node_j_t = edge_node_stack_t[:, :, :, self.dim_nodes : 2 * self.dim_nodes] edge_h_t = edge_node_stack_t[:, :, :, 2 * self.dim_nodes :] if include_nodes: node_i_t = self.node_norm_layer(node_i_t, mask_ij_t) node_j_t = self.node_norm_layer(node_j_t, mask_ij_t) if include_edges: edge_h_t = self.edge_norm_layer(edge_h_t, mask_ij_t) edge_node_stack_t = torch.cat([node_i_t, node_j_t, edge_h_t], -1) return edge_node_stack_t def _update_nodes( self, node_h, node_h_norm, edge_h_norm, edge_idx, mask_i=None, mask_ij=None ): """Update nodes given adjacent nodes and edges""" # Compute messages at each ij edge_node_stack = pack_edges(node_h_norm, edge_h_norm, edge_idx) messages = self.message_mlp(edge_node_stack) if mask_ij is not None: messages = messages * mask_ij.unsqueeze(-1) # Aggregate messages if self.attentional: message = self.attend(node_h_norm, messages, mask_ij) else: message = messages.sum(2) / self.scale node_stack = torch.cat([node_h_norm, message], -1) # Update nodes given aggregated messages node_h_out = node_h + self.update_mlp(node_stack) if mask_i is not None: node_h_out = node_h_out * mask_i.unsqueeze(-1) return node_h_out def _update_nodes_t( self, t, node_h, node_h_norm_t, edge_h_norm_t, edge_idx_t, mask_i_t=None, mask_ij_t=None, ): """Update nodes at index t given adjacent nodes and edges""" # Compute messages at each ij edge_node_stack_t = mask_ij_t.unsqueeze(-1) * pack_edges_step( t, node_h, edge_h_norm_t, edge_idx_t ) # Apply normalization of gathered tensors edge_node_stack_t = self._normalize_t( edge_node_stack_t, mask_ij_t, include_edges=False ) messages_t = self.message_mlp(edge_node_stack_t) if mask_ij_t is not None: messages_t = messages_t * mask_ij_t.unsqueeze(-1) # Aggregate messages if self.attentional: message_t = self.attend(node_h_norm_t, messages_t, mask_ij_t) else: message_t = messages_t.sum(2) / self.scale node_stack_t = torch.cat([node_h_norm_t, message_t], -1) # Update nodes given aggregated messages node_h_t = node_h[:, t, :].unsqueeze(1) node_h_out_t = node_h_t + self.update_mlp(node_stack_t) if mask_i_t is not None: node_h_out_t = node_h_out_t * mask_i_t.unsqueeze(-1) return node_h_out_t def _update_edges(self, edge_h, node_h_out, edge_h_norm, edge_idx, mask_ij): """Update edges given adjacent nodes and edges""" edge_node_stack = pack_edges(node_h_out, edge_h_norm, edge_idx) edge_h_out = edge_h + self.edge_mlp(edge_node_stack) if mask_ij is not None: edge_h_out = edge_h_out * mask_ij.unsqueeze(-1) return edge_h_out def _update_edges_t( self, t, edge_h_t, node_h_out, edge_h_t_norm, edge_idx_t, mask_ij_t ): """Update edges given adjacent nodes and edges""" edge_node_stack_t = pack_edges_step(t, node_h_out, edge_h_t_norm, edge_idx_t) edge_h_out_t = edge_h_t + self.edge_mlp(edge_node_stack_t) if mask_ij_t is not None: edge_h_out_t = edge_h_out_t * mask_ij_t.unsqueeze(-1) return edge_h_out_t def forward( self, node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: Optional[torch.Tensor] = None, mask_ij: Optional[torch.Tensor] = None, ): node_h_norm, edge_h_norm = self._normalize(node_h, edge_h, mask_i, mask_ij) if mask_i is not None: mask_i = (mask_i != 0).type(torch.float32) if mask_ij is not None: mask_ij = (mask_ij != 0).type(torch.float32) node_h_out = self._update_nodes( node_h, node_h_norm, edge_h_norm, edge_idx, mask_i, mask_ij ) edge_h_out = None if self.edge_update: edge_h_out = self._update_edges( edge_h, node_h_out, edge_h_norm, edge_idx, mask_ij ) return node_h_out, edge_h_out def step( self, t: int, node_h: torch.Tensor, node_h_out: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: Optional[torch.Tensor] = None, mask_ij: Optional[torch.Tensor] = None, ): """Compute update for a single node index `t`. This function can be useful for sequential computation of graph updates, for example with autoregressive architectures. Args: t (int): Index of node dimension to update node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, dim_nodes)`. node_h_out (torch.Tensor): Cached outputs of preceding steps with shape `(num_batch, num_nodes, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i (tensor, optional): Node mask with shape `(num_batch, num_nodes)` mask_ij (tensor, optional): Edge mask with shape `(num_batch, num_nodes, num_neighbors)` Resturns: node_h_t (torch.Tensor): Updated node features with shape `(num_batch, 1, dim_nodes)`. edge_h_t (torch.Tensor): Updated edge features with shape `(num_batch, 1, num_neighbors, dim_nodes)`. """ node_h_t = node_h[:, t, :].unsqueeze(1) edge_h_t = edge_h[:, t, :, :].unsqueeze(1) edge_idx_t = edge_idx[:, t, :].unsqueeze(1) mask_i_t = mask_i[:, t].unsqueeze(1) mask_ij_t = mask_ij[:, t, :].unsqueeze(1) """ For a single step we need to apply the normalization both at node t and also for all of the neighborhood tensors that feed in at t. """ node_h_t_norm, edge_h_t_norm = self._normalize( node_h_t, edge_h_t, mask_i_t, mask_ij_t ) node_h_t = self._update_nodes_t( t, node_h, node_h_t_norm, edge_h_t_norm, edge_idx_t, mask_i_t, mask_ij_t ) if self.edge_update: node_h_out = node_h_out.scatter( 1, (t * torch.ones_like(node_h_t)).long(), node_h_t ) edge_h_t = self._update_edges_t( t, edge_h_t, node_h_out, edge_h_t_norm, edge_idx_t, mask_ij_t ) return node_h_t, edge_h_t ## 单纯进行线性变换:Equivariance class MLP(nn.Module): """Multilayer perceptron with variable input, hidden, and output dims. Args: dim_in (int): Feature dimension of input tensor. dim_hidden (int or None): Feature dimension of intermediate layers. Defaults to matching output dimension. dim_out (int or None): Feature dimension of output tensor. Defaults to matching input dimension. num_layers_hidden (int): Number of hidden MLP layers. activation (str): MLP nonlinearity. `'relu'`: Rectified linear unit. `'softplus'`: Softplus. dropout (float): Dropout rate. Default is 0. Inputs: h (torch.Tensor): Input tensor with shape `(..., dim_in)` Outputs: h (torch.Tensor): Input tensor with shape `(..., dim_in)` """ def __init__( self, dim_in: int, dim_hidden: Optional[int] = None, dim_out: Optional[int] = None, num_layers_hidden: int = 1, activation: str = "relu", dropout: float = 0.0, ): super(MLP, self).__init__() # Default is dimension preserving dim_out = dim_out if dim_out is not None else dim_in dim_hidden = dim_hidden if dim_hidden is not None else dim_out nonlinearites = {"relu": nn.ReLU, "softplus": nn.Softplus} activation_func = nonlinearites[activation] if num_layers_hidden == 0: layers = [nn.Linear(dim_in, dim_out)] else: layers = [] for i in range(num_layers_hidden): d_1 = dim_in if i == 0 else dim_hidden layers = layers + [ nn.Linear(d_1, dim_hidden), activation_func(), nn.Dropout(dropout), ] layers = layers + [nn.Linear(dim_hidden, dim_out)] self.layers = nn.Sequential(*layers) def forward(self, h: torch.Tensor) -> torch.Tensor: return self.layers(h) def collect_neighbors(node_h: torch.Tensor, edge_idx: torch.Tensor) -> torch.Tensor: """Collect neighbor node features as edge features. For each node i, collect the embeddings of neighbors {j in N(i)} as edge features neighbor_ij. Args: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, num_features)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. Returns: neighbor_h (torch.Tensor): Edge features containing neighbor node information with shape `(num_batch, num_nodes, num_neighbors, num_features)`. """ num_batch, num_nodes, num_neighbors = edge_idx.shape num_features = node_h.shape[2] # Flatten for the gather operation then reform the full tensor idx_flat = edge_idx.reshape([num_batch, num_nodes * num_neighbors, 1]) idx_flat = idx_flat.expand(-1, -1, num_features) neighbor_h = torch.gather(node_h, 1, idx_flat) neighbor_h = neighbor_h.reshape((num_batch, num_nodes, num_neighbors, num_features)) return neighbor_h def collect_edges( edge_h_dense: torch.Tensor, edge_idx: torch.LongTensor ) -> torch.Tensor: """Collect sparse edge features from a dense pairwise tensor. Args: edge_h_dense (torch.Tensor): Dense edges features with shape `(num_batch, num_nodes, num_nodes, num_features)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. Returns: edge_h (torch.Tensor): Edge features with shape (num_batch, num_nodes, num_neighbors, num_features)`. """ gather_idx = edge_idx.unsqueeze(-1).expand(-1, -1, -1, edge_h_dense.size(-1)) edge_h = torch.gather(edge_h_dense, 2, gather_idx) return edge_h def collect_edges_transpose( edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_ij: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Collect edge embeddings of reversed (transposed) edges in-place. Args: edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, num_features_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)` Returns: edge_h_transpose (torch.Tensor): Edge features of transpose with shape `(num_batch, num_nodes, num_neighbors, num_features_edges)`. mask_ji (torch.Tensor): Mask indicating presence of reversed edge with shape `(num_batch, num_nodes, num_neighbors)`. """ num_batch, num_residues, num_k, num_features = list(edge_h.size()) # Get indices of reverse edges ij_to_ji, mask_ji = transpose_edge_idx(edge_idx, mask_ij) # Gather features at reverse edges edge_h_flat = edge_h.reshape(num_batch, num_residues * num_k, -1) ij_to_ji = ij_to_ji.unsqueeze(-1).expand(-1, -1, num_features) edge_h_transpose = torch.gather(edge_h_flat, 1, ij_to_ji) edge_h_transpose = edge_h_transpose.reshape( num_batch, num_residues, num_k, num_features ) edge_h_transpose = mask_ji.unsqueeze(-1) * edge_h_transpose return edge_h_transpose, mask_ji def scatter_edges(edge_h: torch.Tensor, edge_idx: torch.LongTensor) -> torch.Tensor: """Scatter sparse edge features into a dense pairwise tensor. Args: edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, num_features_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. Returns: edge_h_dense (torch.Tensor): Dense edge features with shape `(batch_size, num_nodes, num_nodes, dimensions)`. """ assert edge_h.dim() == 4 assert edge_idx.dim() == 3 bs, nres, _, dim = edge_h.size() edge_indices = edge_idx.unsqueeze(-1).repeat(1, 1, 1, dim) result = torch.zeros( size=(bs, nres, nres, dim), dtype=edge_h.dtype, device=edge_h.device, ) return result.scatter(dim=2, index=edge_indices, src=edge_h) def pack_edges( node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor ) -> torch.Tensor: """Pack nodes and edge features into edge features. Expands each edge_ij by packing node i, node j, and edge ij into {node,node,edge}_ij. Args: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, num_features_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, num_features_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. Returns: edge_packed (torch.Tensor): Concatenated node and edge features with shape (num_batch, num_nodes, num_neighbors, num_features_nodes + 2*num_features_edges)`. """ num_neighbors = edge_h.shape[2] node_i = node_h.unsqueeze(2).expand(-1, -1, num_neighbors, -1) node_j = collect_neighbors(node_h, edge_idx) edge_packed = torch.cat([node_i, node_j, edge_h], -1) return edge_packed def pack_edges_step( t: int, node_h: torch.Tensor, edge_h_t: torch.Tensor, edge_idx_t: torch.LongTensor ) -> torch.Tensor: """Pack node and edge features into edge features for a single node index t. Expands each edge_ij by packing node i, node j, and edge ij into {node,node,edge}_ij. Args: t (int): Node index to decode. node_h (torch.Tensor): Node features at all positions with shape `(num_batch, num_nodes, num_features_nodes)`. edge_h_t (torch.Tensor): Edge features at index `t` with shape `(num_batch, 1, num_neighbors, num_features_edges)`. edge_idx_t (torch.LongTensor): Edge indices at index `t` for neighbors with shape `(num_batch, 1, num_neighbors)`. Returns: edge_packed (torch.Tensor): Concatenated node and edge features for index `t` with shape (num_batch, 1, num_neighbors, num_features_nodes + 2*num_features_edges)`. """ num_nodes_i = node_h.shape[1] num_neighbors = edge_h_t.shape[2] node_h_t = node_h[:, t, :].unsqueeze(1) node_i = node_h_t.unsqueeze(2).expand(-1, -1, num_neighbors, -1) node_j = collect_neighbors(node_h, edge_idx_t) edge_packed = torch.cat([node_i, node_j, edge_h_t], -1) return edge_packed def transpose_edge_idx( edge_idx: torch.LongTensor, mask_ij: torch.Tensor ) -> Tuple[torch.LongTensor, torch.Tensor]: """Collect edge indices of reverse edges in-place at each edge. The tensor `edge_idx` stores a directed graph topology as a tensor of neighbor indices, where an element `edge_idx[b,i,k]` corresponds to the node index of neighbor `k` of node `i` in batch member `b`. This function takes a directed graph topology and returns an index tensor that maps, in-place, to the reversed edges (if they exist). The indices correspond to the contracted dimension of `edge_index` when it is viewed as `(num_batch, num_nodes * num_neighbors)`. These indices can be used in conjunction with `torch.gather` to collect edge embeddings of `j->i` at `i->j`. See `collect_edges_transpose` for an example. For reverse `j->i` edges that do not exist in the directed graph, the function also returns a binary mask `mask_ji` indicating which edges have both `i->j` and `j->i` present in the graph. Args: edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)` Returns: ij_to_ji (torch.LongTensor): Flat indices for indexing ji in-place at ij with shape `(num_batch, num_nodes * num_neighbors)`. mask_ji (torch.Tensor): Mask indicating presence of reversed edge with shape `(num_batch, num_nodes, num_neighbors)`. """ num_batch, num_residues, num_k = list(edge_idx.size()) # 1. Collect neighbors of neighbors edge_idx_flat = edge_idx.reshape([num_batch, num_residues * num_k, 1]).expand( -1, -1, num_k ) edge_idx_neighbors = torch.gather(edge_idx, 1, edge_idx_flat) # (b,i,j,k) gives the kth neighbor of the jth neighbor of i edge_idx_neighbors = edge_idx_neighbors.reshape( [num_batch, num_residues, num_k, num_k] ) # 2. Determine which k at j maps back to i (if it exists) residue_i = torch.arange(num_residues, device=edge_idx.device).reshape( (1, -1, 1, 1) ) edge_idx_match = (edge_idx_neighbors == residue_i).type(torch.float32) return_mask, return_idx = torch.max(edge_idx_match, -1) # 3. Build flat indices ij_to_ji = edge_idx * num_k + return_idx ij_to_ji = ij_to_ji.reshape(num_batch, -1) # 4. Transpose mask mask_ji = torch.gather(mask_ij.reshape(num_batch, -1), -1, ij_to_ji) mask_ji = mask_ji.reshape(num_batch, num_residues, num_k) mask_ji = mask_ij * return_mask * mask_ji return ij_to_ji, mask_ji def permute_tensor( tensor: torch.Tensor, dim: int, permute_idx: torch.LongTensor ) -> torch.Tensor: """Permute a tensor along a dimension given a permutation vector. Args: tensor (torch.Tensor): Input tensor with shape `([batch_dims], permutation_length, [content_dims])`. dim (int): Dimension to permute along. permute_idx (torch.LongTensor): Permutation index tensor with shape `([batch_dims], permutation_length)`. Returns: tensor_permute (torch.Tensor): Permuted node features with shape `([batch_dims], permutation_length, [content_dims])`. """ # Resolve absolute dimension dim = range(len(list(tensor.shape)))[dim] # Flatten content dimensions shape = list(tensor.shape) batch_dims, permute_length = shape[:dim], shape[dim] tensor_flat = tensor.reshape(batch_dims + [permute_length] + [-1]) # Exap content dimensions permute_idx_expand = permute_idx.unsqueeze(-1).expand(tensor_flat.shape) tensor_permute_flat = torch.gather(tensor_flat, dim, permute_idx_expand) tensor_permute = tensor_permute_flat.reshape(tensor.shape) return tensor_permute def permute_graph_embeddings( node_h: torch.Tensor, edge_h: torch.Tensor, edge_idx: torch.LongTensor, mask_i: torch.Tensor, mask_ij: torch.Tensor, permute_idx: torch.LongTensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor, torch.Tensor, torch.Tensor]: """Permute graph embeddings given a permutation vector. Args: node_h (torch.Tensor): Node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h (torch.Tensor): Edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i (tensor, optional): Node mask with shape `(num_batch, num_nodes)` mask_ij (tensor, optional): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. permute_idx (torch.LongTensor): Permutation vector with shape `(num_batch, num_nodes)`. Returns: node_h_permute (torch.Tensor): Permuted node features with shape `(num_batch, num_nodes, dim_nodes)`. edge_h_permute (torch.Tensor): Permuted edge features with shape `(num_batch, num_nodes, num_neighbors, dim_edges)`. edge_idx_permute (torch.LongTensor): Permuted edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_i_permute (tensor, optional): Permuted node mask with shape `(num_batch, num_nodes)` mask_ij_permute (tensor, optional): Permuted edge mask with shape `(num_batch, num_nodes, num_neighbors)`. """ # Permuting one-dimensional objects is straightforward gathering node_h_permute = permute_tensor(node_h, 1, permute_idx) edge_h_permute = permute_tensor(edge_h, 1, permute_idx) mask_i_permute = permute_tensor(mask_i, 1, permute_idx) mask_ij_permute = permute_tensor(mask_ij, 1, permute_idx) """ For edge_idx, there are two-dimensions set each edge idx that previously pointed to j to now point to the new location of j which is p^(-1)[j] edge^(p)[i,k] = p^(-1)[edge[p(i),k]] """ # First, permute on the i dimension edge_idx_permute_1 = permute_tensor(edge_idx, 1, permute_idx) # Second, permute on the j dimension by using the inverse permute_idx_inverse = torch.argsort(permute_idx, dim=-1) edge_idx_1_flat = edge_idx_permute_1.reshape([edge_idx.shape[0], -1]) edge_idx_permute_flat = torch.gather(permute_idx_inverse, 1, edge_idx_1_flat) edge_idx_permute = edge_idx_permute_flat.reshape(edge_idx.shape) return ( node_h_permute, edge_h_permute, edge_idx_permute, mask_i_permute, mask_ij_permute, ) def edge_mask_causal(edge_idx: torch.LongTensor, mask_ij: torch.Tensor) -> torch.Tensor: """Make an edge mask causal with mask_ij = 0 for j >= i. Args: edge_idx (torch.LongTensor): Edge indices for neighbors with shape `(num_batch, num_nodes, num_neighbors)`. mask_ij (torch.Tensor): Edge mask with shape `(num_batch, num_nodes, num_neighbors)`. Returns: mask_ij_causal (torch.Tensor): Causal edge mask with shape `(num_batch, num_nodes, num_neighbors)`. """ idx = torch.arange(edge_idx.size(1), device=edge_idx.device) idx_expand = idx.reshape([1, -1, 1]) mask_ij_causal = (edge_idx < idx_expand).float() * mask_ij return mask_ij_causal class MaskedNorm(nn.Module): """Masked normalization layer. Args: dim (int): Dimensionality of the normalization. Can be 1 for 1D normalization along dimension 1 or 2 for 2D normalization along dimensions 1 and 2. num_features (int): Channel dimension; only needed if `affine` is True. affine (bool): If True, inclde a learnable affine transformation post-normalization. Default is False. norm (str): Type of normalization, can be `instance`, `layer`, or `transformer`. eps (float): Small number for numerical stability. Inputs: data (torch.Tensor): Input tensor with shape `(num_batch, num_nodes, num_channels)` (1D) or `(num_batch, num_nodes, num_nodes, num_channels)` (2D). mask (torch.Tensor): Mask tensor with shape `(num_batch, num_nodes)` (1D) or `(num_batch, num_nodes, num_nodes)` (2D). Outputs: norm_data (torch.Tensor): Mask-normalized tensor with shape `(num_batch, num_nodes, num_channels)` (1D) or `(num_batch, num_nodes, num_nodes, num_channels)` (2D). """ def __init__( self, dim: int, num_features: int = -1, affine: bool = False, norm: str = "instance", eps: float = 1e-5, ): super(MaskedNorm, self).__init__() self.norm_type = norm self.dim = dim self.norm = norm + str(dim) self.affine = affine self.eps = eps # Dimension to sum if self.norm == "instance1": self.sum_dims = [1] elif self.norm == "layer1": self.sum_dims = [1, 2] elif self.norm == "transformer1": self.sum_dims = [-1] elif self.norm == "instance2": self.sum_dims = [1, 2] elif self.norm == "layer2": self.sum_dims = [1, 2, 3] elif self.norm == "transformer2": self.sum_dims = [-1] else: raise NotImplementedError # Number of features, only required if affine self.num_features = num_features # Affine transformation is a linear layer on the C channel if self.affine: self.weights = nn.Parameter(torch.rand(self.num_features)) self.bias = nn.Parameter(torch.zeros(self.num_features)) def forward( self, data: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: # Add optional trailing singleton dimension and expand if necessary if mask is not None: if len(mask.shape) == len(data.shape) - 1: mask = mask.unsqueeze(-1) if data.shape != mask.shape: mask = mask.expand(data.shape) # Input shape is Batch, Channel, Dim1, (dim2 if 2d) dims = self.sum_dims if (mask is None) or (self.norm_type == "transformer"): mask_mean = data.mean(dim=dims, keepdim=True) mask_std = torch.sqrt( (((data - mask_mean)).pow(2)).mean(dim=dims, keepdim=True) + self.eps ) # Norm norm_data = (data - mask_mean) / mask_std else: # Zeroes vector to sum all mask data norm_data = torch.zeros_like(data).to(data.device).type(data.dtype) for mask_id in mask.unique(): # Skip zero, since real mask if mask_id == 0: continue # Transform mask to temp mask that match mask id tmask = (mask == mask_id).type(torch.float32) # Sum mask for mean mask_sum = tmask.sum(dim=dims, keepdim=True) # Data is tmask, so that mean is only for unmasked pos mask_mean = (data * tmask).sum(dim=dims, keepdim=True) / mask_sum mask_std = torch.sqrt( (((data - mask_mean) * tmask).pow(2)).sum(dim=dims, keepdim=True) / mask_sum + self.eps ) # Calculate temp norm, apply mask tnorm = ((data - mask_mean) / mask_std) * tmask # Sometime mask is empty, so generate nan that are conversted to 0 tnorm[tnorm != tnorm] = 0 # Add to init zero norm data norm_data += tnorm # Apply affine if self.affine: norm_data = norm_data * self.weights + self.bias # If mask, apply mask if mask is not None: norm_data = norm_data * (mask != 0).type(data.dtype) return norm_data