noether.core.schemas.modules.attention

Classes

AttentionConfig	Configuration for an attention module.
DotProductAttentionConfig	Configuration for the Dot Product attention module.
TransolverAttentionConfig	Configuration for the Transolver attention module.
TransolverPlusPlusAttentionConfig	Configuration for the Transolver++ attention module.
IrregularNatAttentionConfig	Configuration for the Irregular Neighbourhood Attention Transformer (NAT) attention module.
PerceiverAttentionConfig	Configuration for the Perceiver attention module.
MultiBranchAnchorAttentionConfig	Configuration for Multi-Branch Anchor Attention module.
CrossAnchorAttentionConfig	Configuration for Cross Anchor Attention module.
JointAnchorAttentionConfig	Configuration for Joint Anchor Attention module.
TokenSpec	Specification for a token type in the attention mechanism.
AttentionPattern	Defines which tokens attend to which other tokens.
MixedAttentionConfig	Configuration for Mixed Attention module.

Module Contents

class noether.core.schemas.modules.attention.AttentionConfig(/, **data)

Bases: pydantic.BaseModel

Configuration for an attention module. Since we can have many different attention implementations, we allow extra fields. such that we can use the same schema for all attention modules.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

model_config

Configuration for an attention module.

hidden_dim: int = None

Dimensionality of the hidden features.

num_heads: int = None

Number of attention heads.

use_rope: bool = None

Whether to use Rotary Positional Embeddings (RoPE).

dropout: float = None

Dropout rate for the attention weights and output projection.

init_weights: noether.core.types.InitWeightsMode = None

Weight initialization strategy.

bias: bool = None

Whether to use bias terms in linear layers.

head_dim: int | None = None

Dimensionality of each attention head.

validate_hidden_dim_and_num_heads()

class noether.core.schemas.modules.attention.DotProductAttentionConfig(/, **data)

Bases: AttentionConfig

Configuration for the Dot Product attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

class noether.core.schemas.modules.attention.TransolverAttentionConfig(/, **data)

Bases: AttentionConfig

Configuration for the Transolver attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

num_slices: int = None

Number of slices to project the input tokens to.

class noether.core.schemas.modules.attention.TransolverPlusPlusAttentionConfig(/, **data)

Bases: TransolverAttentionConfig

Configuration for the Transolver++ attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

use_overparameterization: bool = None

Whether to use overparameterization for the slice projection.

use_adaptive_temperature: bool = None

Whether to use an adaptive temperature for the slice selection.

temperature_activation: Literal['sigmoid', 'softplus', 'exp'] | None = None

Activation function for the adaptive temperature.

use_gumbel_softmax: bool = None

Whether to use Gumbel-Softmax for the slice selection.

class noether.core.schemas.modules.attention.IrregularNatAttentionConfig(/, **data)

Bases: AttentionConfig

Configuration for the Irregular Neighbourhood Attention Transformer (NAT) attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

input_dim: int = None

Dimensionality of the input features.

radius: float = None

Radius for the radius graph.

max_degree: int = None

Maximum number of neighbors per point.

relpos_mlp_hidden_dim: int = None

Hidden dimensionality of the relative position bias MLP.

relpos_mlp_dropout: float = None

Dropout rate for the relative position bias MLP.

class noether.core.schemas.modules.attention.PerceiverAttentionConfig(/, **data)

Bases: AttentionConfig

Configuration for the Perceiver attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

kv_dim: int | None = None

Dimensionality of the key/value features. If None, use hidden_dim.

set_kv_dim()

class noether.core.schemas.modules.attention.MultiBranchAnchorAttentionConfig(/, **data)

Bases: AttentionConfig

Configuration for Multi-Branch Anchor Attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

branches: list[str] = None

anchor_suffix: str = None

class noether.core.schemas.modules.attention.CrossAnchorAttentionConfig(/, **data)

Bases: MultiBranchAnchorAttentionConfig

Configuration for Cross Anchor Attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

class noether.core.schemas.modules.attention.JointAnchorAttentionConfig(/, **data)

Bases: MultiBranchAnchorAttentionConfig

Configuration for Joint Anchor Attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

class noether.core.schemas.modules.attention.TokenSpec(/, **data)

Bases: pydantic.BaseModel

Specification for a token type in the attention mechanism.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

name: Literal['surface_anchors', 'volume_anchors', 'surface_queries', 'volume_queries']

size: int = None

classmethod from_dict(token_dict)

Create TokenSpec from dictionary with single key-value pair.

Parameters:

token_dict (dict[str, int])

Return type:

TokenSpec

to_dict()

Convert TokenSpec to dictionary.

Return type:

dict[str, int]

class noether.core.schemas.modules.attention.AttentionPattern(/, **data)

Bases: pydantic.BaseModel

Defines which tokens attend to which other tokens.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)

query_tokens: collections.abc.Sequence[str]

key_value_tokens: collections.abc.Sequence[str]

class noether.core.schemas.modules.attention.MixedAttentionConfig(/, **data)

Bases: DotProductAttentionConfig

Configuration for Mixed Attention module.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

data (Any)