codingutf8 2018 the google flax team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license bertinputsdocstring r args inputids obj torch longtensor of shape obj 0 indices of input sequence tokens in the vocabulary indices can be obtained using class transformers berttokenizer see meth transformers pretrainedtokenizer encode and meth transformers pretrainedtokenizer call for details what are input ids glossary htmlinputids attentionmask obj torch floattensor of shape obj 0 optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary htmlattentionmask tokentypeids obj torch longtensor of shape obj 0 optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossary htmltokentypeids positionids obj torch longtensor of shape obj 0 optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config maxpositionembeddings 1 what are position ids glossary htmlpositionids headmask obj torch floattensor of shape obj numheads or obj numlayers numheads optional mask to nullify selected heads of the selfattention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputsembeds obj torch floattensor of shape obj 0 hiddensize optional optionally instead of passing obj inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert obj inputids indices into associated vectors than the model s internal embedding lookup matrix outputattentions obj bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates obj bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict obj bool optional whether or not to return a class transformers fileutils modeloutput instead of a plain tuple layer normalization https arxiv orgabs1607 06450 operates on the last axis of the input data also e g nn relu this can be disabled since the scaling will be done by the next layer applies layer normalization on the input it normalizes the activations of the layer for each given example in a batch independently rather than across a batch like batch normalization i e applies a transformation that maintains the mean activation within each example close to 0 and the activation standard deviation close to 1 args x the inputs returns normalized inputs the same shape as inputs specify a new class for doing the embedding stuff as flax s one use embedding for the parameter name and pytorch use weight construct the embeddings from word position and tokentype embeddings vocabsize int hiddensize int typevocabsize int maxlength int nn compact def callself inputids tokentypeids positionids attentionmask embed wemb flaxperformerembeddingself vocabsize self hiddensize namewordembeddings jnp atleast2dinputids astypei4 pemb flaxperformerembeddingself maxlength self hiddensize namepositionembeddings jnp atleast2dpositionids astypei4 temb flaxperformerembeddingself typevocabsize self hiddensize nametokentypeembeddings jnp atleast2dtokentypeids astypei4 sum all embeddings summedemb wemb jnp broadcasttopemb wemb shape temb layer norm layernorm flaxperformerlayernormnamelayernormsummedemb return layernorm class flaxperformerattentionnn module numheads int headsize int nn compact def callself hiddenstate attentionmask singleheaddim self headsize self numheads fastsoftmaxattention makefastsoftmaxattentionqkvdimsingleheaddim selfatt nn attention selfattention numheadsself numheads qkvfeaturesself headsize nameself attentionfnfastsoftmaxattention hiddenstate attentionmask layernorm flaxperformerlayernormnamelayernormselfatt hiddenstate return layernorm class flaxperformerintermediatenn module outputsize int hiddenact str gelu nn compact def callself hiddenstate todo add act2fn reference to change activation function dense nn densefeaturesself outputsize namedensehiddenstate return act2fnself hiddenactdense class flaxperformeroutputnn module nn compact def callself intermediateoutput attentionoutput hiddenstate nn denseattentionoutput shape1 namedenseintermediateoutput hiddenstate flaxperformerlayernormnamelayernormhiddenstate attentionoutput return hiddenstate class flaxperformerlayernn module numheads int headsize int intermediatesize int hiddenact str gelu nn compact def callself hiddenstate attentionmask attention flaxperformerattentionself numheads self headsize nameattention hiddenstate attentionmask intermediate flaxperformerintermediate self intermediatesize nameintermediate hiddenactself hiddenact attention output flaxperformeroutputnameoutputintermediate attention return output class flaxperformerlayercollectionnn module numlayers int numheads int headsize int intermediatesize int hiddenact str gelu nn compact def callself inputs attentionmask assert self numlayers 0 fnumlayers should be 1 got self numlayers initialize input output inputi inputs forward over all encoders for i in rangeself numlayers layer flaxperformerlayer self numheads self headsize self intermediatesize hiddenactself hiddenact namefi inputi layerinputi attentionmask return inputi class flaxperformerencodernn module numlayers int numheads int headsize int intermediatesize int hiddenact str gelu nn compact def callself hiddenstate attentionmask layer flaxperformerlayercollection self numlayers self numheads self headsize self intermediatesize namelayer hiddenactself hiddenact hiddenstate attentionmask return layer class flaxperformerpoolernn module nn compact def callself hiddenstate clstoken hiddenstate 0 out nn densehiddenstate shape1 namedenseclstoken return jax lax tanhout class flaxperformermodulenn module vocabsize int hiddensize int typevocabsize int maxlength int numencoderlayers int numheads int headsize int intermediatesize int hiddenact str gelu addpoolinglayer bool true nn compact def callself inputids tokentypeids positionids attentionmask embedding embeddings flaxperformerembeddings self vocabsize self hiddensize self typevocabsize self maxlength nameembeddings inputids tokentypeids positionids attentionmask n stacked encoding layers encoder flaxperformerencoder self numencoderlayers self numheads self headsize self intermediatesize hiddenactself hiddenact nameencoder embeddings attentionmask if not self addpoolinglayer return encoder pooled flaxperformerpoolernamepoolerencoder return encoder pooled addstartdocstrings the bare bert model transformer outputting raw hiddenstates without any specific head on top bertstartdocstring class flaxperformermodelflaxbertpretrainedmodel modelclass flaxperformermodule configclass bertconfig basemodelprefix bert staticmethod def convertfrompytorchptstate dict config bertconfig dict jaxstate dictptstate need to change some parameters name to match flax names so that we don t have to fork any layer for key tensor in ptstate items key parts keyparts setkey split every dense layer has kernel parameters instead of weight if dense weight in key del jaxstatekey key key replaceweight kernel jaxstatekey tensor selfattention needs also to replace weight by kernel if query key value keyparts flax selfattention decomposes the heads numhead size numheads if bias in key jaxstatekey tensor reshapeconfig numattentionheads 1 elif weight del jaxstatekey key key replaceweight kernel tensor tensor reshapeconfig numattentionheads 1 config hiddensize transpose2 0 1 jaxstatekey tensor selfattention output is not a separate layer remove one nesting if attention output dense in key del jaxstatekey key key replaceattention output dense attention self out jaxstatekey tensor selfattention output is not a separate layer remove nesting on layer norm if attention output layernorm in key del jaxstatekey key key replaceattention output layernorm attention layernorm jaxstatekey tensor there are some transposed parameters w r t their pytorch counterpart if intermediate dense kernel in key or output dense kernel in key jaxstatekey tensor t self attention output projection needs to be transposed if out kernel in key jaxstatekey tensor reshapeconfig hiddensize config numattentionheads 1 transpose 1 2 0 pooler needs to transpose its kernel if pooler dense kernel in key jaxstatekey tensor t handle layernorm conversion if layernorm in key del jaxstatekey replace layernorm by layernorm newkey key replacelayernorm layernorm if weight in key newkey newkey replaceweight gamma elif bias in key newkey newkey replacebias beta jaxstatenewkey tensor return jaxstate def init self config bertconfig inputshape tuple 1 1 seed int 0 dtype jnp dtype jnp float32 kwargs module flaxperformermodule vocabsizeconfig vocabsize hiddensizeconfig hiddensize typevocabsizeconfig typevocabsize maxlengthconfig maxpositionembeddings numencoderlayersconfig numhiddenlayers numheadsconfig numattentionheads headsizeconfig hiddensize intermediatesizeconfig intermediatesize dropoutrateconfig hiddendropoutprob hiddenactconfig hiddenact super initconfig module inputshapeinputshape seedseed dtypedtype property def moduleself nn module return self module def call self inputids tokentypeidsnone positionidsnone dropoutrng prngkey none attentionmasknone inputids attentionmask tokentypeids positionids self checkinputs inputids attentionmask tokentypeids positionids handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng return self module apply params self params jnp arrayinputids dtypei4 jnp arraytokentypeids dtypei4 jnp arraypositionids dtypei4 jnp arrayattentionmask dtypei4 rngrngs class flaxperformerformaskedlmflaxbertpretrainedmodel def init self config bertconfig inputshape tuple 1 1 seed int 0 dtype jnp dtype jnp float32 kwargs module flaxperformerformaskedlmmodule vocabsizeconfig vocabsize typevocabsizeconfig typevocabsize hiddensizeconfig hiddensize intermediatesizeconfig intermediatesize headsizeconfig hiddensize numheadsconfig numattentionheads numencoderlayersconfig numhiddenlayers maxlengthconfig maxpositionembeddings hiddenactconfig hiddenact kwargs super initconfig module inputshapeinputshape seedseed dtypedtype def call self inputids attentionmasknone tokentypeidsnone positionidsnone params dict none train bool false dropoutrng prngkey none inputids attentionmask tokentypeids positionids self checkinputs inputids attentionmask tokentypeids positionids handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng return self module apply params params or self params jnp arrayinputids dtypei4 jnp arrayattentionmask dtypei4 jnp arraytokentypeids dtypei4 jnp arraypositionids dtypei4 not train rngsrngs class flaxperformerformaskedlmmodulenn module vocabsize int hiddensize int intermediatesize int headsize int numheads int numencoderlayers int typevocabsize int maxlength int hiddenact str dropoutrate float 0 0 dtype jnp dtype jnp float32 nn compact def call self inputids attentionmasknone tokentypeidsnone positionidsnone deterministic bool true model encoder flaxperformermodule vocabsizeself vocabsize hiddensizeself hiddensize typevocabsizeself typevocabsize maxlengthself maxlength numencoderlayersself numencoderlayers numheadsself numheads headsizeself hiddensize intermediatesizeself intermediatesize hiddenactself hiddenact addpoolinglayerfalse namebert inputids attentionmask tokentypeids positionids compute the prediction scores encoder nn dropoutrateself dropoutrateencoder deterministicdeterministic logits flaxbertonlymlmhead vocabsizeself vocabsize hiddenactself hiddenact namecls dtypeself dtype encoder return logits coding utf 8 2018 the google flax team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license this model inherits from class transformers pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn module https pytorch org docs stable nn html torch nn module __ subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config class transformers bertconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the meth transformers pretrainedmodel from_pretrained method to load the model weights args input_ids obj torch longtensor of shape obj 0 indices of input sequence tokens in the vocabulary indices can be obtained using class transformers berttokenizer see meth transformers pretrainedtokenizer encode and meth transformers pretrainedtokenizer __call__ for details what are input ids glossary html input ids __ attention_mask obj torch floattensor of shape obj 0 optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary html attention mask __ token_type_ids obj torch longtensor of shape obj 0 optional segment token indices to indicate first and second portions of the inputs indices are selected in 0 1 0 corresponds to a sentence a token 1 corresponds to a sentence b token what are token type ids glossary html token type ids _ position_ids obj torch longtensor of shape obj 0 optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 what are position ids glossary html position ids _ head_mask obj torch floattensor of shape obj num_heads or obj num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputs_embeds obj torch floattensor of shape obj 0 hidden_size optional optionally instead of passing obj input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert obj input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions obj bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states obj bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict obj bool optional whether or not to return a class transformers file_utils modeloutput instead of a plain tuple layer normalization https arxiv org abs 1607 06450 operates on the last axis of the input data the dtype of the computation if true bias beta is added if true multiply by scale gamma when the next layer is linear also e g nn relu this can be disabled since the scaling will be done by the next layer applies layer normalization on the input it normalizes the activations of the layer for each given example in a batch independently rather than across a batch like batch normalization i e applies a transformation that maintains the mean activation within each example close to 0 and the activation standard deviation close to 1 args x the inputs returns normalized inputs the same shape as inputs specify a new class for doing the embedding stuff as flax s one use embedding for the parameter name and pytorch use weight construct the embeddings from word position and token_type embeddings embed sum all embeddings layer norm todo add act2fn reference to change activation function stores n bertlayer s initialize input output forward over all encoders embedding n stacked encoding layers the model can behave as an encoder with only self attention as well as a decoder in which case a layer of cross attention is added between the self attention layers following the architecture described in attention is all you need https arxiv org abs 1706 03762 __ by ashish vaswani noam shazeer niki parmar jakob uszkoreit llion jones aidan n gomez lukasz kaiser and illia polosukhin need to change some parameters name to match flax names so that we don t have to fork any layer key parts every dense layer has kernel parameters instead of weight selfattention needs also to replace weight by kernel flax selfattention decomposes the heads num_head size num_heads selfattention output is not a separate layer remove one nesting selfattention output is not a separate layer remove nesting on layer norm there are some transposed parameters w r t their pytorch counterpart self attention output projection needs to be transposed pooler needs to transpose its kernel handle layernorm conversion replace layernorm by layer_norm handle any prng if needed handle any prng if needed model compute the prediction scores
from typing import Callable, Dict, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
from jax.random import PRNGKey
from modeling_flax_performer_utils import make_fast_softmax_attention

from transformers.file_utils import add_start_docstrings
from transformers.modeling_flax_utils import ACT2FN
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.bert.modeling_flax_bert import FlaxBertOnlyMLMHead, FlaxBertPreTrainedModel
from transformers.utils import logging


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "BertConfig"
_TOKENIZER_FOR_DOC = "BertTokenizer"

BERT_START_DOCSTRING = r

BERT_INPUTS_DOCSTRING = r


class FlaxPerformerLayerNorm(nn.Module):
    

    epsilon: float = 1e-6
    dtype: jnp.dtype = jnp.float32  
    bias: bool = True  
    scale: bool = True  
    
    
    bias_init: jnp.ndarray = nn.initializers.zeros
    scale_init: jnp.ndarray = nn.initializers.ones

    @nn.compact
    def __call__(self, x):
        
        features = x.shape[-1]
        mean = jnp.mean(x, axis=-1, keepdims=True)
        mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
        var = mean2 - jax.lax.square(mean)
        mul = jax.lax.rsqrt(var + self.epsilon)
        if self.scale:
            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)), self.dtype)
        y = (x - mean) * mul
        if self.bias:
            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)), self.dtype)
        return y


class FlaxPerformerEmbedding(nn.Module):
    

    vocab_size: int
    hidden_size: int
    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)

    @nn.compact
    def __call__(self, inputs):
        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))
        return jnp.take(embedding, inputs, axis=0)


class FlaxPerformerEmbeddings(nn.Module):
    

    vocab_size: int
    hidden_size: int
    type_vocab_size: int
    max_length: int

    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
        
        w_emb = FlaxPerformerEmbedding(self.vocab_size, self.hidden_size, name="word_embeddings")(
            jnp.atleast_2d(input_ids.astype("i4"))
        )
        p_emb = FlaxPerformerEmbedding(self.max_length, self.hidden_size, name="position_embeddings")(
            jnp.atleast_2d(position_ids.astype("i4"))
        )
        t_emb = FlaxPerformerEmbedding(self.type_vocab_size, self.hidden_size, name="token_type_embeddings")(
            jnp.atleast_2d(token_type_ids.astype("i4"))
        )

        
        summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb

        
        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(summed_emb)

        return layer_norm


class FlaxPerformerAttention(nn.Module):
    num_heads: int
    head_size: int

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        single_head_dim = self.head_size // self.num_heads
        fast_softmax_attention = make_fast_softmax_attention(qkv_dim=single_head_dim)
        self_att = nn.attention.SelfAttention(
            num_heads=self.num_heads, qkv_features=self.head_size, name="self", attention_fn=fast_softmax_attention
        )(hidden_state, attention_mask)

        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(self_att + hidden_state)
        return layer_norm


class FlaxPerformerIntermediate(nn.Module):
    output_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state):
        
        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
        return ACT2FN[self.hidden_act](dense)


class FlaxPerformerOutput(nn.Module):
    @nn.compact
    def __call__(self, intermediate_output, attention_output):
        hidden_state = nn.Dense(attention_output.shape[-1], name="dense")(intermediate_output)
        hidden_state = FlaxPerformerLayerNorm(name="layer_norm")(hidden_state + attention_output)
        return hidden_state


class FlaxPerformerLayer(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        attention = FlaxPerformerAttention(self.num_heads, self.head_size, name="attention")(
            hidden_state, attention_mask
        )
        intermediate = FlaxPerformerIntermediate(
            self.intermediate_size, name="intermediate", hidden_act=self.hidden_act
        )(attention)
        output = FlaxPerformerOutput(name="output")(intermediate, attention)

        return output


class FlaxPerformerLayerCollection(nn.Module):
    

    num_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, inputs, attention_mask):
        assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"

        
        input_i = inputs

        
        for i in range(self.num_layers):
            layer = FlaxPerformerLayer(
                self.num_heads, self.head_size, self.intermediate_size, hidden_act=self.hidden_act, name=f"{i}"
            )
            input_i = layer(input_i, attention_mask)
        return input_i


class FlaxPerformerEncoder(nn.Module):
    num_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"

    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        layer = FlaxPerformerLayerCollection(
            self.num_layers,
            self.num_heads,
            self.head_size,
            self.intermediate_size,
            name="layer",
            hidden_act=self.hidden_act,
        )(hidden_state, attention_mask)
        return layer


class FlaxPerformerPooler(nn.Module):
    @nn.compact
    def __call__(self, hidden_state):
        cls_token = hidden_state[:, 0]
        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
        return jax.lax.tanh(out)


class FlaxPerformerModule(nn.Module):
    vocab_size: int
    hidden_size: int
    type_vocab_size: int
    max_length: int
    num_encoder_layers: int
    num_heads: int
    head_size: int
    intermediate_size: int
    hidden_act: str = "gelu"
    add_pooling_layer: bool = True

    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
        
        embeddings = FlaxPerformerEmbeddings(
            self.vocab_size, self.hidden_size, self.type_vocab_size, self.max_length, name="embeddings"
        )(input_ids, token_type_ids, position_ids, attention_mask)

        
        encoder = FlaxPerformerEncoder(
            self.num_encoder_layers,
            self.num_heads,
            self.head_size,
            self.intermediate_size,
            hidden_act=self.hidden_act,
            name="encoder",
        )(embeddings, attention_mask)

        if not self.add_pooling_layer:
            return encoder

        pooled = FlaxPerformerPooler(name="pooler")(encoder)
        return encoder, pooled


@add_start_docstrings(
    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
    BERT_START_DOCSTRING,
)
class FlaxPerformerModel(FlaxBertPreTrainedModel):
    

    model_class = FlaxPerformerModule
    config_class = BertConfig
    base_model_prefix = "bert"

    @staticmethod
    def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
        jax_state = dict(pt_state)

        
        for key, tensor in pt_state.items():
            
            key_parts = set(key.split("."))

            
            if "dense.weight" in key:
                del jax_state[key]
                key = key.replace("weight", "kernel")
                jax_state[key] = tensor

            
            if {"query", "key", "value"} & key_parts:
                
                if "bias" in key:
                    jax_state[key] = tensor.reshape((config.num_attention_heads, -1))
                elif "weight":
                    del jax_state[key]
                    key = key.replace("weight", "kernel")
                    tensor = tensor.reshape((config.num_attention_heads, -1, config.hidden_size)).transpose((2, 0, 1))
                    jax_state[key] = tensor

            
            if "attention.output.dense" in key:
                del jax_state[key]
                key = key.replace("attention.output.dense", "attention.self.out")
                jax_state[key] = tensor

            
            if "attention.output.LayerNorm" in key:
                del jax_state[key]
                key = key.replace("attention.output.LayerNorm", "attention.LayerNorm")
                jax_state[key] = tensor

            
            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:
                jax_state[key] = tensor.T

            
            if "out.kernel" in key:
                jax_state[key] = tensor.reshape((config.hidden_size, config.num_attention_heads, -1)).transpose(
                    1, 2, 0
                )

            
            if "pooler.dense.kernel" in key:
                jax_state[key] = tensor.T

            
            if "LayerNorm" in key:
                del jax_state[key]

                
                new_key = key.replace("LayerNorm", "layer_norm")

                if "weight" in key:
                    new_key = new_key.replace("weight", "gamma")
                elif "bias" in key:
                    new_key = new_key.replace("bias", "beta")

                jax_state[new_key] = tensor

        return jax_state

    def __init__(
        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
    ):
        module = FlaxPerformerModule(
            vocab_size=config.vocab_size,
            hidden_size=config.hidden_size,
            type_vocab_size=config.type_vocab_size,
            max_length=config.max_position_embeddings,
            num_encoder_layers=config.num_hidden_layers,
            num_heads=config.num_attention_heads,
            head_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            dropout_rate=config.hidden_dropout_prob,
            hidden_act=config.hidden_act,
        )

        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

    @property
    def module(self) -> nn.Module:
        return self._module

    def __call__(
        self, input_ids, token_type_ids=None, position_ids=None, dropout_rng: PRNGKey = None, attention_mask=None
    ):
        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
            input_ids, attention_mask, token_type_ids, position_ids
        )

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        return self.module.apply(
            {"params": self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            rng=rngs,
        )


class FlaxPerformerForMaskedLM(FlaxBertPreTrainedModel):
    def __init__(
        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
    ):
        module = FlaxPerformerForMaskedLMModule(
            vocab_size=config.vocab_size,
            type_vocab_size=config.type_vocab_size,
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            head_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_encoder_layers=config.num_hidden_layers,
            max_length=config.max_position_embeddings,
            hidden_act=config.hidden_act,
            **kwargs,
        )

        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

    def __call__(
        self,
        input_ids,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        params: dict = None,
        train: bool = False,
        dropout_rng: PRNGKey = None,
    ):
        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
            input_ids, attention_mask, token_type_ids, position_ids
        )

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            not train,
            rngs=rngs,
        )


class FlaxPerformerForMaskedLMModule(nn.Module):
    vocab_size: int
    hidden_size: int
    intermediate_size: int
    head_size: int
    num_heads: int
    num_encoder_layers: int
    type_vocab_size: int
    max_length: int
    hidden_act: str
    dropout_rate: float = 0.0
    dtype: jnp.dtype = jnp.float32

    @nn.compact
    def __call__(
        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
    ):
        
        encoder = FlaxPerformerModule(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            type_vocab_size=self.type_vocab_size,
            max_length=self.max_length,
            num_encoder_layers=self.num_encoder_layers,
            num_heads=self.num_heads,
            head_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            hidden_act=self.hidden_act,
            add_pooling_layer=False,
            name="bert",
        )(input_ids, attention_mask, token_type_ids, position_ids)

        
        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
        logits = FlaxBertOnlyMLMHead(
            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="cls", dtype=self.dtype
        )(encoder)

        return (logits,)