noether.core.schemas.modules¶

Submodules¶

Classes¶

`AttentionConfig`	Configuration for an attention module.
`CrossAnchorAttentionConfig`	Configuration for Cross Anchor Attention module.
`DotProductAttentionConfig`	Configuration for the Dot Product attention module.
`IrregularNatAttentionConfig`	Configuration for the Irregular Neighbourhood Attention Transformer (NAT) attention module.
`JointAnchorAttentionConfig`	Configuration for Joint Anchor Attention module.
`MixedAttentionConfig`	Configuration for Mixed Attention module.
`MultiBranchAnchorAttentionConfig`	Configuration for Multi-Branch Anchor Attention module.
`PerceiverAttentionConfig`	Configuration for the Perceiver attention module.
`TransolverAttentionConfig`	Configuration for the Transolver attention module.
`TransolverPlusPlusAttentionConfig`	Configuration for the Transolver++ attention module.
`PerceiverBlockConfig`	Configuration for the PerceiverBlock module.
`TransformerBlockConfig`	Configuration for a transformer block.
`DeepPerceiverDecoderConfig`	Configuration for the DeepPerceiverDecoder module.
`SupernodePoolingConfig`
`ContinuousSincosEmbeddingConfig`	Configuration for Continuous Sine-Cosine Embedding layer.
`LayerScaleConfig`	Configuration for Layer Scale module.
`LinearProjectionConfig`	Configuration for a LinearProjection layer.
`RopeFrequencyConfig`	Configuration for RoPE frequency settings.
`UnquantizedDropPathConfig`	Configuration for the UnquantizedDropPath layer.
`MLPConfig`
`UpActDownMLPConfig`

Package Contents¶

class noether.core.schemas.modules.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.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.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.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.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.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)

class noether.core.schemas.modules.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.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.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.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.PerceiverBlockConfig(/, **data)¶

Bases: TransformerBlockConfig

Configuration for the PerceiverBlock 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 and value representations. Defaults to None. If None, hidden_dim is used.

set_kv_dim()¶

Set kv_dim to hidden_dim if not provided.

Return type:: PerceiverBlockConfig

perceiver_attention_config()¶

Return type:: noether.core.schemas.modules.attention.PerceiverAttentionConfig

modulation_linear_projection_config()¶

Return type:: noether.core.schemas.modules.layers.LinearProjectionConfig | None

class noether.core.schemas.modules.TransformerBlockConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for a transformer block.

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)

hidden_dim: int = None¶: Hidden Dimension of the transformer block.

num_heads: int = None¶: Number of attention heads.

mlp_hidden_dim: int | None = None¶: Hidden dimension of the MLP layer. If set to None, the mlp_hidden dim is set to hidden_dim * mlp_expansion_factor in the TransformerConfig. If both are None, an error is raised.

mlp_expansion_factor: int | None = None¶: Expansion factor for the MLP hidden dimension relative to the hidden dimension. If ‘mlp_hidden_dim’ is not set, this factor is used to compute it as hidden_dim * mlp_expansion_factor.

drop_path: float = None¶: Probability to drop the attention or MLP module. Defaults to 0.0.

attention_constructor: Literal['dot_product', 'perceiver', 'transolver', 'transolver_plusplus'] = 'dot_product'¶: Constructor of the attention module. Defaults to ‘dot_product’.

layerscale: float | None = None¶: Init scale value to scale layer activations. Defaults to None.

condition_dim: int | None = None¶: Dimension of the conditioning vector. If none, no conditioning is applied. If provided, the transformer block will turn into a Diffusion Transformer (DiT) block.

bias: bool = None¶: Whether to use biases in norm/projections. Defaults to True.

eps: float = None¶: Epsilon Value for the layer nornalization. Defaults to 1e-6.

init_weights: noether.core.types.InitWeightsMode = None¶: Initialization method for the weight matrices of the network. Defaults to “truncnormal002

use_rope: bool = None¶: Whether to use Rotary Positional Embeddings (RoPE).

attention_arguments: dict¶: Additional arguments for the attention module that are only needed for a specific attention implementation.

set_mlp_hidden_dim()¶

linear_projection_config()¶

Return type:: noether.core.schemas.modules.layers.LinearProjectionConfig

layerscale_config()¶

Return type:: noether.core.schemas.modules.layers.LayerScaleConfig

drop_path_config()¶

Return type:: noether.core.schemas.modules.layers.UnquantizedDropPathConfig

modulation_linear_projection_config()¶

Return type:: LinearProjectionConfig | None

up_act_down_mlp_config()¶

Return type:: noether.core.schemas.modules.mlp.UpActDownMLPConfig

class noether.core.schemas.modules.DeepPerceiverDecoderConfig(/, **data)¶

Bases: noether.core.schemas.mixins.InjectSharedFieldFromParentMixin, pydantic.BaseModel

Configuration for the DeepPerceiverDecoder 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)

perceiver_block_config: Annotated[noether.core.schemas.modules.blocks.PerceiverBlockConfig, noether.core.schemas.mixins.Shared] = None¶: Configuration for the Perceiver blocks used in the decoder.

depth: int = None¶: Number of deep perceiver decoder layers (i.e., depth of the network). Defaults to 1.

input_dim: int = None¶: Input dimension for the query positions.

class noether.core.schemas.modules.SupernodePoolingConfig(/, **data)¶

Bases: pydantic.BaseModel

Parameters:: data (Any)

hidden_dim: int = None¶: Hidden dimension for positional embeddings, messages and the resulting output vector.

input_dim: int = None¶: Number of positional dimension (e.g., input_dim=2 for a 2D position, input_dim=3 for a 3D position)

radius: float | None = None¶: Radius around each supernode. From points within this radius, messages are passed to the supernode.

k: int | None = None¶: Number of neighbors for each supernode. From the k-NN points, messages are passed to the supernode.

max_degree: int = None¶: Maximum degree of the radius graph. Defaults to 32.

spool_pos_mode: Literal['abspos', 'relpos', 'absrelpos'] = None¶

absolute space (“abspos”), relative space (“relpos”) or both (“absrelpos”).

Type:: Type of position embedding

init_weights: noether.core.types.InitWeightsMode = None¶: Weight initialization of linear layers. Defaults to “truncnormal002”.

readd_supernode_pos: bool = None¶: If true, the absolute positional encoding of the supernode is concatenated to the supernode vector after message passing and linearly projected back to hidden_dim. Defaults to True.

aggregation: Literal['mean', 'sum'] = None¶: Aggregation for message passing (“mean” or “sum”).

message_mode: Literal['mlp', 'linear', 'identity'] = None¶: How messages are created. “mlp” (2 layer MLP), “linear” (nn.Linear), “identity” (nn.Identity). Defaults to “mlp”.

input_features_dim: int | None = None¶: Number of input features per point. None will fall back to a version without features. Defaults to None, which means no input features.

validate_radius_and_k()¶

class noether.core.schemas.modules.ContinuousSincosEmbeddingConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for Continuous Sine-Cosine Embedding layer.

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)

hidden_dim: int = None¶: Dimensionality of the output embedding.

input_dim: int = None¶: Dimensionality of the input coordinates.

max_wavelength: int = None¶: Maximum wavelength for the sine-cosine embeddings.

class noether.core.schemas.modules.LayerScaleConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for Layer Scale 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)

hidden_dim: int = None¶: Number of dimensions of the input tensor to be scaled.

init_values: float | None = None¶: Initial gamme scale value. Defaults to 1e-5.

class noether.core.schemas.modules.LinearProjectionConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for a LinearProjection layer.

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¶: Input dimension of the linear projection.

output_dim: int = None¶: Output dimension of the linear projection.

ndim: None | int = None¶: Number of dimensions of the input domain. Either None (Linear projection), 1D (sequence), 2D, or 3D. Defaults to None.

bias: bool = None¶: If true, use bias term in the linear projection. Defaults to True.

optional: bool = None¶: If true and input_dim==output_dim (i.e., there is no up/down projection), then the identity mapping is used. Defaults to False.

init_weights: noether.core.types.InitWeightsMode = None¶: Initialization method of the weights of the MLP. Options are ‘torch’ (i.e., similar to the module) or ‘truncnormal002’, or ‘zero’. Defaults to ‘torch’.

validate_ndim()¶

Validate the ndim field to ensure it is either None, 1, 2, or 3.

Return type:: Self

class noether.core.schemas.modules.RopeFrequencyConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for RoPE frequency settings.

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)

hidden_dim: int = None¶: Dimensionality of frequencies (in transformers this should be the head dimension).

input_dim: int = None¶: Dimensionality of the coordinates (e.g., 2 for 2D coordinates, 3 for 3D coordinates).

max_wavelength: int = None¶

10000.0

Type:: Theta parameter for the transformer sine/cosine embedding. Default

implementation: Literal['real', 'complex'] = None¶: “real” -> basic implementation using real coordinates (this is slow and only here for backward compatibility). “complex” -> fast implementation of rotation via complex multiplication. Default: “real”.

class noether.core.schemas.modules.UnquantizedDropPathConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for the UnquantizedDropPath layer.

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)

drop_prob: float = None¶: Probability of dropping a path during training.

scale_by_keep: bool = None¶: Up-scales activations during training by 1 - drop_prob to avoid train-test mismatch. Defaults to True.

class noether.core.schemas.modules.MLPConfig(/, **data)¶

Bases: pydantic.BaseModel

Parameters:: data (Any)

input_dim: int = None¶: Input dimension of the MLP.

output_dim: int = None¶: Output dimension of the MLP.

hidden_dim: int = None¶: Hidden dimension for each layer.

num_layers: int = None¶: Number of hidden layers in the MLP. If 0, the MLP is a two linear layer MLP from input_dim, hidden_dim, activation to output_dim.

activation: Literal['RELU', 'GELU', 'SIGMOID', 'TANH', 'LEAKY_RELU', 'SOFTPLUS', 'ELU', 'SILU'] = 'GELU'¶: Activation function to use between layers.

init_weights: noether.core.types.InitWeightsMode = 'truncnormal002'¶: Weight initialization method.

class noether.core.schemas.modules.UpActDownMLPConfig(/, **data)¶

Bases: pydantic.BaseModel

Parameters:: data (Any)

input_dim: int = None¶: Input dimension of the MLP.

hidden_dim: int = None¶: Hidden dimension of the MLP.

bias: bool = None¶: Whether to use bias in the MLP.

init_weights: noether.core.types.InitWeightsMode = None¶: Initialization method of the weights of the MLP. Options are “torch” (i.e., similar to the module) or ‘truncnormal002’. Defaults to ‘truncnormal002’.

check_dims()¶

Validator to check that hidden_dim is greater than input_dim.

Raises:: ValueError – raised if hidden_dim is not greater than input_dim.
Return type:: UpActDownMLPConfig