src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraTiedDense/__call__
class FlaxElectraTiedDense:    def __call__(self, x, kernel):
        x = jnp.asarray(x, self.dtype)
        kernel = jnp.asarray(kernel, self.dtype)
        y = lax.dot_general(
            x,
            kernel,
            (((x.ndim - 1,), (0,)), ((), ())),
            precision=self.precision,
        )
        bias = jnp.asarray(self.bias, self.dtype)
        return y + bias