src/transformers/models/distilbert/modeling_flax_distilbert.py/FlaxDistilBertLMDecoder/__call__
class FlaxDistilBertLMDecoder:    def __call__(self, inputs, kernel):
        y = lax.dot_general(inputs, kernel, (((inputs.ndim - 1,), (0,)), ((), ())))
        y = y + self.bias
        return y