Class FlexBertEncoderBase

src/bert_layers/layers.py:596–610

A FlexBERT base class for type hints.

class FlexBertEncoderBase(nn.Module):
    """A FlexBERT base class for type hints."""

    layers: nn.ModuleList

    def _init_weights(self, reset_params: bool = False):
        if hasattr(self, "layers"):
            for layer in self.layers:
                layer._init_weights(reset_params=reset_params)

    def reset_parameters(self):
        self._init_weights(reset_params=True)

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError("This is a base class and should not be used directly.")


class FlexBertUnpadEncoder(FlexBertEncoderBase):