Method _create_layer_parameter_specs

axlearn/common/bert.py:158–167

(self)

        )

    def _create_layer_parameter_specs(self) -> dict[str, ParameterSpec]:
        cfg = self.config
        return dict(
            # Output layer weights are shared with encoder, but bias is a separate parameter.
            output_bias=ParameterSpec(
                shape=[cfg.num_classes],
                mesh_axes=(None,),
                initializer=constant_initializer(0.0),
            ),
        )

    def forward(self, input_batch: NestedTensor) -> Tensor:
        """Produces prediction scores from BERT output features.