noether.modeling.modules.layers

Submodules

Classes

ContinuousSincosEmbed

Embedding layer for continuous coordinates using sine and cosine functions.

UnquantizedDropPath

Unquantized drop path (Stochastic Depth, https://arxiv.org/abs/1603.09382) per sample. Unquantized means

LayerScale

LayerScale module scales the input tensor by a learnable parameter gamma.

LinearProjection

LinearProjection is a linear projection layer that can be used for 1D, 2D, and 3D data.

RopeFrequency

Creates frequencies for rotary embeddings (RoPE) from https://arxiv.org/abs/2104.09864 for variable positions.

ScalarsConditioner

Embeds num_scalars scalars into a single conditioning vector via first encoding every scalar with

TransformerBatchNorm

Wrapper around torch.nn.BatchNorm1d that considers all tokens of a single sample as the full batch.

Package Contents

class noether.modeling.modules.layers.ContinuousSincosEmbed(config)

Bases: torch.nn.Module

Embedding layer for continuous coordinates using sine and cosine functions. The original implementation from the Attenion is All You Need paper, deals with descrete 1D cordinates (i.e., a sequence). Howerver, this implementation is able to deal with 2D and 3D coordinate systems as well.

Parameters:

config (noether.core.schemas.modules.layers.ContinuousSincosEmbeddingConfig) – Configuration for the ContinuousSincosEmbed module. See ContinuousSincosEmbeddingConfig for the available options.

omega: torch.Tensor
padding_tensor: torch.Tensor
hidden_dim
input_dim
ndim_padding
sincos_padding
max_wavelength
padding
forward(coords)

Forward method of the ContinuousSincosEmbed layer.

Parameters:

coords (torch.Tensor) – Tensor of coordinates. The shape of the tensor should be [batch size, number of points, coordinate dimension] or [number of points, coordinate dimension].

Raises:

NotImplementedError – Only supports sparse (i.e. [number of points, coordinate dimension]) or dense (i.e. [batch size, number of points, coordinate dimension]) coordinates systems.

Returns:

Tensor with embedded coordinates.

Return type:

torch.Tensor

class noether.modeling.modules.layers.UnquantizedDropPath(config)

Bases: torch.nn.Module

Unquantized drop path (Stochastic Depth, https://arxiv.org/abs/1603.09382) per sample. Unquantized means that dropped paths are still calculated. Number of dropped paths is fully stochastic, i.e., it can happen that not a single path is dropped or that all paths are dropped. In a quantized drop path, the same amount of paths are dropped in each forward pass, resulting in large speedups with high drop_prob values. See https://arxiv.org/abs/2212.04884 for more discussion. UnquantizedDropPath does not provide any speedup, consider using a quantized version if large drop_prob values are used.

Adapted from https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py#L150

Initialize the UnquantizedDropPath module.

Parameters:

config (noether.core.schemas.modules.layers.UnquantizedDropPathConfig) – Configuration for the UnquantizedDropPath module. See UnquantizedDropPathConfig for the available options.

drop_prob
scale_by_keep
property keep_prob

Return the keep probability. I.e. the probability to keep a path, which is 1 - drop_prob.

Returns:

Float value of the keep probability.

forward(x)

Forward function of the UnquantizedDropPath module.

Parameters:

x (torch.Tensor) – Tensor to apply the drop path. Shape: (batch_size, …).

Returns:

(batch_size, …). If drop_prob is 0, the input tensor is returned. If drop_prob is 1, a tensor with zeros is returned.

Return type:

Tensor with drop path applied. Shape

extra_repr()

Extra representation of the UnquantizedDropPath module.

Returns:

Return a string representation of the module.

class noether.modeling.modules.layers.LayerScale(config)

Bases: torch.nn.Module

LayerScale module scales the input tensor by a learnable parameter gamma.

Initialize the LayerScale module. :param config: Configuration for the LayerScale module. See LayerScaleConfig for details.

Parameters:

config (noether.core.schemas.modules.layers.LayerScaleConfig)

forward(x)

Forward function of the LayerScale module.

Parameters:

x (torch.Tensor) – Input tensor to be scaled.

Returns:

Tensor scaled by the gamma parameter.

Return type:

torch.Tensor

class noether.modeling.modules.layers.LinearProjection(config)

Bases: torch.nn.Module

LinearProjection is a linear projection layer that can be used for 1D, 2D, and 3D data.

Parameters:

config (noether.core.schemas.modules.layers.LinearProjectionConfig) – The configuration of the LinearProjection. See LinearProjectionConfig for available options.

Raises:

NotImplementedError – raises not implemented error if the number of dimensions of the input domain is bigger than 4.

project: torch.nn.Linear | torch.nn.Conv1d | torch.nn.Conv2d | torch.nn.Conv3d | torch.nn.Identity
init_weights
reset_parameters()
Reset the parameters of the MLP with a specific initialization. Options are “torch” (i.e., default) or

“truncnormal002”.

Raises:

NotImplementedError – raised if the specified initialization is not implemented.

Return type:

None

forward(x)

Forward function of the LinearProjection.

Parameters:

x (torch.Tensor) – Input tensor to the LinearProjection.

Returns:

Output tensor from the LinearProjection.

Return type:

torch.Tensor

class noether.modeling.modules.layers.RopeFrequency(config)

Bases: torch.nn.Module

Creates frequencies for rotary embeddings (RoPE) from https://arxiv.org/abs/2104.09864 for variable positions.

Parameters:

config (noether.core.schemas.modules.layers.RopeFrequencyConfig) – Configuration for RoPE frequency settings. See RopeFrequencyConfig for available options.

omega: torch.Tensor
hidden_dim
input_dim
implementation
ndim_padding
sincos_padding
max_wavelength
padding
forward(coords)
Parameters:

coords (torch.Tensor) – coordinates to create RoPE frequencies for. Expected shape is (…, input_dim).

Return type:

torch.Tensor | tuple[torch.Tensor, Ellipsis]

class noether.modeling.modules.layers.ScalarsConditioner(config)

Bases: torch.nn.Module

Embeds num_scalars scalars into a single conditioning vector via first encoding every scalar with sine-cosine embeddings followed by a mlp (per scalar). These vectors are then concatenated and projected down to condition_dim with an MLP.

Parameters:

config (noether.core.schemas.modules.layers.scalar_conditioner.ScalarsConditionerConfig) – configuration for the ScalarsConditioner. See ScalarsConditionerConfig for available options.

hidden_dim
num_scalars
condition_dim
embed
mlps
shared_mlp
forward(*args, **kwargs)

Embeds scalars into a single conditioning vector. Scalars can be passed as *args or as **kwargs. It is recommended to use kwargs to avoid bugs that originate from passing scalars in a different order at two locations in the code. Recommended usage: condition = conditioner(geometry_angle=75.3, friction_angle=24.6) :param *args: Scalars in tensor representation (batch_size,) or (batch_size, 1). :param **kwargs: Scalars in tensor representation (batch_size,) or (batch_size, 1).

Returns:

Conditioning vector with shape (batch_size, condition_dim)

Parameters:
Return type:

torch.Tensor

Example: .. code-block:: python

conditioner = ScalarsConditioner(
ScalarsConditionerConfig(

hidden_dim=64, num_scalars=2, condition_dim=128, init_weights=”gaussian”,

)

) geometry_angle = torch.tensor([75.3, 80.1]) # shape (batch_size,) friction_angle = torch.tensor([24.6, 30.2]) # shape (batch_size,) condition = conditioner(

geometry_angle=geometry_angle, friction_angle=friction_angle

) # shape (batch_size, condition_dim)

class noether.modeling.modules.layers.TransformerBatchNorm(num_features, eps=1e-05, elementwise_affine=True, bias=True)

Bases: torch.nn.Module

Wrapper around torch.nn.BatchNorm1d that considers all tokens of a single sample as the full batch. Additionally remaps affine to elementwise_affine and supports disabling bias to comply with the torch.nn.LayerNorm interface. Does not use any nn.BatchNorm1d modules to avoid errors with nn.SyncBatchnorm.

Parameters:
num_features
eps = 1e-05
elementwise_affine = True
forward(x)

BatchNorm1d where all tokens of a single sample correspond to a full batch.

Parameters:

x (torch.Tensor) – Tensor of shape (batch_size, seqlen, dim).

Returns:

Normalized x of shape (batch_size, seqlen, dim).

Return type:

torch.Tensor