noether.modeling.modules.attention.transolver_plusplus¶
Classes¶
Transolver++ Attention module as implemented in https://github.com/thuml/Transolver_plus/blob/main/models/Transolver_plus.py |
Module Contents¶
- class noether.modeling.modules.attention.transolver_plusplus.TransolverPlusPlusAttention(config)¶
Bases:
torch.nn.ModuleTransolver++ Attention module as implemented in https://github.com/thuml/Transolver_plus/blob/main/models/Transolver_plus.py
Initialize the TransolverPlusPlusAttention module.
- Parameters:
config (noether.core.schemas.modules.AttentionConfig) – Configuration object for the attention module.
- dim_head¶
- num_heads = None¶
- scale¶
- softmax¶
- dropout = None¶
- bias¶
- proj_temperature¶
- in_project_x¶
- in_project_slice¶
- qkv¶
- to_out¶
- forward(x, attn_mask=None)¶
Forward pass of the Transolver attention module.
- Parameters:
x (torch.Tensor) – Input tensor with shape (batch_size, seqlen, hidden_dim).
attn_mask (torch.Tensor | None) – Attention mask tensor with shape (batch_size). Defaults to None.
- Returns:
Tensor after applying the transolver attention mechanism.