(cls)
| 110 | |
| 111 | @classmethod |
| 112 | def default_config(cls): |
| 113 | cfg = super().default_config() |
| 114 | cfg.dtype = jnp.float32 # pylint: disable=no-member |
| 115 | cfg.layer = Linear.default_config().set( |
| 116 | input_dim=_EXAMPLE_SHAPE[-1], |
| 117 | output_dim=_EXAMPLE_SHAPE[-1], |
| 118 | bias=False, |
| 119 | param_partition_spec=("model", None), |
| 120 | ) |
| 121 | cfg.name = cls.__name__ |
| 122 | cfg.param_init = param_init.ConstantInitializer.default_config().set(value=1.0) |
| 123 | return cfg |
| 124 | |
| 125 | def __init__(self, cfg: BaseModel.Config, *, parent: Optional[Module]): |
| 126 | super().__init__(cfg, parent=parent) |
nothing calls this directly
no test coverage detected