noether.modeling.modules.layers.layer_scale

Classes

LayerScaleConfig

Configuration for Layer Scale module.

LayerScale

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

Module Contents

class noether.modeling.modules.layers.layer_scale.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.modeling.modules.layers.layer_scale.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 (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