src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraTiedDense/setup
class FlaxElectraTiedDense:    def setup(self):
        bias = self.param("bias", self.bias_init, (self.embedding_size,))
        self.bias = jnp.asarray(bias, dtype=self.dtype)