Class Config

axlearn/common/bert.py:67–71

    @config_class
    class Config(BaseLayer.Config):
        input_dim: Required[int] = REQUIRED  # Input dimension.
        linear: InstantiableConfig = Linear.default_config()
        dropout: InstantiableConfig = Dropout.default_config()
        activation: str = "jnp.tanh"

    def __init__(self, cfg: Config, *, parent: Optional[Module]):
        super().__init__(cfg, parent=parent)