pytagi.nn.rms_norm

pytagi.nn.rms_norm#

Classes#

RMSNorm

Implements Root Mean Square Layer Normalization (RMSNorm) by normalizing

Module Contents#

class pytagi.nn.rms_norm.RMSNorm(normalized_shape: List[int], eps: float = 1e-06)[source]#

Bases: pytagi.nn.base_layer.BaseLayer

Implements Root Mean Square Layer Normalization (RMSNorm) by normalizing the inputs using the RMS statistic. It inherits from BaseLayer.

Initializes the RMSNorm layer.

Parameters:

  • normalized_shape – The shape of the input to normalize over (e.g., the size of the feature dimension). Expected to be a list of integers.

  • eps – A small value added to the denominator for numerical stability to prevent division by zero. Defaults to 1e-6.

get_layer_info() → str[source]#

Retrieves a descriptive string containing information about the layer’s configuration from the C++ backend.

get_layer_name() → str[source]#

Retrieves the name of the layer (e.g., ‘RMSNorm’) from the C++ backend.

init_weight_bias()[source]#

Initializes the layer’s internal parameters, specifically the learnable scale (gamma). This task is delegated to the C++ backend.

Contents