src/transformers/models/distilbert/modeling_flax_distilbert.py/FlaxDistilBertLMDecoder/setup
class FlaxDistilBertLMDecoder:    def setup(self):
        self.bias = self.param("bias", self.bias_init, (self.config.vocab_size,))