codingutf8 2021 the fairseq s and 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 flax blenderbot model import math import random from functools import partial from typing import callable optional tuple import flax linen as nn import jax import jax numpy as jnp from flax core frozendict import frozendict freeze unfreeze from flax linen import combinemasks makecausalmask from flax linen attention import dotproductattentionweights from flax traverseutil import flattendict unflattendict from jax import lax from jax random import prngkey from modelingflaxoutputs import flaxbasemodeloutput flaxbasemodeloutputwithpastandcrossattentions flaxcausallmoutputwithcrossattentions flaxseq2seqlmoutput flaxseq2seqmodeloutput from modelingflaxutils import act2fn flaxpretrainedmodel appendcallsampledocstring appendreplacereturndocstrings overwritecalldocstring from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationblenderbot import blenderbotconfig logger logging getloggername configfordoc blenderbotconfig checkpointfordoc facebookblenderbot400mdistill blenderbotstartdocstring r this model inherits from flaxpretrainedmodel 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 flax linen flax nn modulehttps flax readthedocs ioenlatestautosummaryflax nn module html subclass use it as a regular flax module and refer to the flax documentation for all matter related to general usage and behavior finally this model supports inherent jax features such as justintime jit compilationhttps jax readthedocs ioenlatestjax htmljustintimecompilationjit automatic differentiationhttps jax readthedocs ioenlatestjax htmlautomaticdifferentiation vectorizationhttps jax readthedocs ioenlatestjax htmlvectorizationvmap parallelizationhttps jax readthedocs ioenlatestjax htmlparallelizationpmap parameters config blenderbotconfig 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 flaxpretrainedmodel frompretrained method to load the model weights blenderbotencodeinputsdocstring r args inputids jnp ndarray of shape batchsize sequencelength indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are input ids glossaryinputids attentionmask jnp ndarray of shape batchsize sequencelength 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 glossaryattentionmask positionids numpy ndarray of shape batchsize sequencelength optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config maxpositionembeddings 1 outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models bart modelingflaxbart shifttokensright def shifttokensrightinputids jnp ndarray padtokenid int decoderstarttokenid int jnp ndarray shiftedinputids jnp zeroslikeinputids shiftedinputids shiftedinputids at 1 setinputids 1 shiftedinputids shiftedinputids at 0 setdecoderstarttokenid shiftedinputids jnp whereshiftedinputids 100 padtokenid shiftedinputids return shiftedinputids copied from transformers models bart modelingflaxbart flaxbartattention with bartblenderbot class flaxblenderbotattentionnn module config blenderbotconfig embeddim int numheads int dropout float 0 0 causal bool false bias bool true dtype jnp dtype jnp float32 the dtype of the computation def setupself none self headdim self embeddim self numheads if self headdim self numheads self embeddim raise valueerror fembeddim must be divisible by numheads got embeddim self embeddim f and numheads self numheads dense partial nn dense self embeddim usebiasself bias dtypeself dtype kernelinitjax nn initializers normalself config initstd self qproj self kproj self vproj dense dense dense self outproj dense self dropoutlayer nn dropoutrateself dropout if self causal self causalmask makecausalmask jnp ones1 self config maxpositionembeddings dtypebool dtypebool def splitheadsself hiddenstates return hiddenstates reshapehiddenstates shape 2 self numheads self headdim def mergeheadsself hiddenstates return hiddenstates reshapehiddenstates shape 2 self embeddim nn compact def concatenatetocacheself key value query attentionmask detect if we re initializing by absence of existing cache data isinitialized self hasvariablecache cachedkey cachedkey self variablecache cachedkey jnp zeros key shape key dtype cachedvalue self variablecache cachedvalue jnp zeros value shape value dtype cacheindex self variablecache cacheindex lambda jnp array0 dtypejnp int32 if isinitialized batchdims maxlength numheads depthperhead cachedkey value shape update key value caches with our new 1d spatial slices curindex cacheindex value indices 0 lenbatchdims curindex 0 0 key lax dynamicupdateslicecachedkey value key indices value lax dynamicupdateslicecachedvalue value value indices cachedkey value key cachedvalue value value numupdatedcachevectors query shape1 cacheindex value cacheindex value numupdatedcachevectors causal mask for cached decoder selfattention our single query position should only attend to those key positions that have already been generated and cached not the remaining zero elements padmask jnp broadcastto jnp arangemaxlength curindex numupdatedcachevectors tuplebatchdims 1 numupdatedcachevectors maxlength attentionmask combinemaskspadmask attentionmask return key value attentionmask def call self hiddenstates jnp ndarray keyvaluestates optionaljnp ndarray none attentionmask optionaljnp ndarray none initcache bool false deterministic bool true tuplejnp ndarray if keyvaluestates are provided this layer is used as a crossattention layer for the decoder get query proj get key value proj crossattentions selfattention handle cache prepare causal attention mask combine masks if needed during fast autoregressive decoding we feed one position at a time and cache the keys and values step by step convert the boolean attention mask to an attention bias attention mask in the form of attention bias copied from transformers models mbart modelingflaxmbart flaxmbartencoderlayer with mbartblenderbot copied from transformers models bart modelingflaxbart flaxbartencoderlayercollection with bartblenderbot add layerdrop see https arxiv orgabs1909 11556 for description copied from transformers models mbart modelingflaxmbart flaxmbartdecoderlayer with mbartblenderbot self attention crossattention block fully connected copied from transformers models bart modelingflaxbart flaxbartdecoderlayercollection with bartblenderbot decoder layers add layerdrop see https arxiv orgabs1909 11556 for description add hidden states from the last decoder layer update the last element in hiddenstates after applying layernorm above embed positions update the last element in hiddenstates after applying layernorm above copied from transformers models bart modelingflaxbart flaxbartmodule with bartblenderbot init input tensors make sure initialization pass will work for flaxblenderbotforsequenceclassificationmodule init input variables to retrieve cache decoderinputids jnp onesbatchsize maxlength dtypei4 decoderattentionmask jnp oneslikedecoderinputids decoderpositionids jnp broadcastto jnp arangejnp atleast2ddecoderinputids shape1 decoderinputids shape def decoderforwardmodule decoderinputids decoderattentionmask decoderpositionids kwargs decodermodule module getdecodermodule return decodermodule decoderinputids decoderattentionmask decoderpositionids kwargs initvariables self module init jax random prngkey0 decoderinputidsdecoderinputids decoderattentionmaskdecoderattentionmask decoderpositionidsdecoderpositionids encoderhiddenstatesencoderoutputs0 initcachetrue methoddecoderforward we only need to call the decoder to init the cache return unfreezeinitvariablescache addstartdocstringsblenderbotencodeinputsdocstring replacereturndocstringsoutputtypeflaxbasemodeloutput configclassblenderbotconfig def encode self inputids jnp ndarray attentionmask optionaljnp ndarray none positionids optionaljnp ndarray none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns example python from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration frompretrainedfacebookblenderbot400mdistill tokenizer autotokenizer frompretrainedfacebookblenderbot400mdistill text my friends are cool but they eat too many carbs inputs tokenizertext maxlength1024 returntensorsjax encoderoutputs model encodeinputs outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict if attentionmask is none attentionmask jnp oneslikeinputids if positionids is none batchsize sequencelength inputids shape positionids jnp broadcasttojnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng def encoderforwardmodule inputids attentionmask positionids kwargs encodemodule module getencodermodule return encodemoduleinputids attentionmask positionids kwargs return self module apply params params or self params inputidsjnp arrayinputids dtypei4 attentionmaskjnp arrayattentionmask dtypei4 positionidsjnp arraypositionids dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs methodencoderforward addstartdocstringsblenderbotdecodeinputsdocstring replacereturndocstrings outputtypeflaxbasemodeloutputwithpastandcrossattentions configclassblenderbotconfig def decode self decoderinputids encoderoutputs encoderattentionmask optionaljnp ndarray none decoderattentionmask optionaljnp ndarray none decoderpositionids optionaljnp ndarray none pastkeyvalues dict none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns example python import jax numpy as jnp from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration frompretrainedfacebookblenderbot400mdistill tokenizer autotokenizer frompretrainedfacebookblenderbot400mdistill text my friends are cool but they eat too many carbs inputs tokenizertext maxlength1024 returntensorsjax encoderoutputs model encodeinputs decoderstarttokenid model config decoderstarttokenid decoderinputids jnp onesinputs inputids shape0 1 dtypei4 decoderstarttokenid outputs model decodedecoderinputids encoderoutputs lastdecoderhiddenstates outputs lasthiddenstate outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict encoderhiddenstates encoderoutputs0 if encoderattentionmask is none batchsize sequencelength encoderhiddenstates shape 2 encoderattentionmask jnp onesbatchsize sequencelength batchsize sequencelength decoderinputids shape if decoderattentionmask is none decoderattentionmask jnp onesbatchsize sequencelength if decoderpositionids is none if pastkeyvalues is not none raise valueerrormake sure to provide decoderpositionids when passing pastkeyvalues decoderpositionids jnp broadcastto jnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng inputs params params or self params if pastkeyvalues are passed then cache is already initialized a private flag initcache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxblenderbotattention module if pastkeyvalues inputscache pastkeyvalues mutable cache else mutable false def decoderforwardmodule decoderinputids decoderattentionmask decoderpositionids kwargs decodermodule module getdecodermodule return decodermodule decoderinputids decoderattentionmask decoderpositionids kwargs outputs self module apply inputs decoderinputidsjnp arraydecoderinputids dtypei4 decoderattentionmaskjnp arraydecoderattentionmask dtypei4 decoderpositionidsjnp arraydecoderpositionids dtypei4 encoderhiddenstatesencoderhiddenstates encoderattentionmaskjnp arrayencoderattentionmask dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs mutablemutable methoddecoderforward add updated cache to model output if pastkeyvalues is not none and returndict outputs past outputs outputspastkeyvalues unfreezepastcache return outputs elif pastkeyvalues is not none and not returndict outputs past outputs outputs outputs 1 unfreezepastcache outputs1 return outputs addstartdocstringstomodelforwardblenderbotinputsdocstring def call self inputids jnp ndarray attentionmask optionaljnp ndarray none decoderinputids optionaljnp ndarray none decoderattentionmask optionaljnp ndarray none positionids optionaljnp ndarray none decoderpositionids optionaljnp ndarray none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict prepare encoder inputs if attentionmask is none attentionmask jnp oneslikeinputids if positionids is none batchsize sequencelength inputids shape positionids jnp broadcasttojnp arangesequencelengthnone batchsize sequencelength prepare decoder inputs if decoderinputids is none decoderinputids shifttokensright inputids self config padtokenid decoderstarttokenidself config decoderstarttokenid if decoderattentionmask is none decoderattentionmask jnp oneslikedecoderinputids if decoderpositionids is none batchsize sequencelength decoderinputids shape decoderpositionids jnp broadcastto jnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs dropout dropoutrng if dropoutrng is not none else return self module apply params params or self params inputidsjnp arrayinputids dtypei4 attentionmaskjnp arrayattentionmask dtypei4 positionidsjnp arraypositionids dtypei4 decoderinputidsjnp arraydecoderinputids dtypei4 decoderattentionmaskjnp arraydecoderattentionmask dtypei4 decoderpositionidsjnp arraydecoderpositionids dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs addstartdocstrings the bare mbart model transformer outputting raw hiddenstates without any specific head on top blenderbotstartdocstring class flaxblenderbotmodelflaxblenderbotpretrainedmodel config blenderbotconfig dtype jnp dtype jnp float32 the dtype of the computation moduleclass flaxblenderbotmodule appendcallsampledocstringflaxblenderbotmodel checkpointfordoc flaxseq2seqmodeloutput configfordoc copied from transformers models bart modelingflaxbart flaxbartforconditionalgenerationmodule with bartblenderbot class flaxblenderbotforconditionalgenerationmodulenn module config blenderbotconfig dtype jnp dtype jnp float32 biasinit callable jnp ndarray jax nn initializers zeros def setupself self model flaxblenderbotmoduleconfigself config dtypeself dtype self lmhead nn dense self model shared numembeddings usebiasfalse dtypeself dtype kernelinitjax nn initializers normalself config initstd self finallogitsbias self paramfinallogitsbias self biasinit 1 self model shared numembeddings def getencodermoduleself return self model encoder def getdecodermoduleself return self model decoder def call self inputids attentionmask decoderinputids decoderattentionmask positionids decoderpositionids outputattentions bool false outputhiddenstates bool false returndict bool true deterministic bool true outputs self model inputidsinputids attentionmaskattentionmask decoderinputidsdecoderinputids decoderattentionmaskdecoderattentionmask positionidspositionids decoderpositionidsdecoderpositionids outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicdeterministic hiddenstates outputs0 if self config tiewordembeddings sharedembedding self model variablesparamssharedembedding lmlogits self lmhead applyparams kernel sharedembedding t hiddenstates else lmlogits self lmheadhiddenstates lmlogits jax lax stopgradientself finallogitsbias astypeself dtype if not returndict output lmlogits outputs1 return output return flaxseq2seqlmoutput logitslmlogits decoderhiddenstatesoutputs decoderhiddenstates decoderattentionsoutputs decoderattentions crossattentionsoutputs crossattentions encoderlasthiddenstateoutputs encoderlasthiddenstate encoderhiddenstatesoutputs encoderhiddenstates encoderattentionsoutputs encoderattentions addstartdocstrings the blenderbot model with a language modeling head can be used for summarization blenderbotstartdocstring class flaxblenderbotforconditionalgenerationflaxblenderbotpretrainedmodel moduleclass flaxblenderbotforconditionalgenerationmodule dtype jnp dtype jnp float32 addstartdocstringsblenderbotdecodeinputsdocstring replacereturndocstringsoutputtypeflaxcausallmoutputwithcrossattentions configclassblenderbotconfig def decode self decoderinputids encoderoutputs encoderattentionmask optionaljnp ndarray none decoderattentionmask optionaljnp ndarray none decoderpositionids optionaljnp ndarray none pastkeyvalues dict none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns example python import jax numpy as jnp from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration frompretrainedfacebookblenderbot400mdistill tokenizer autotokenizer frompretrainedfacebookblenderbot400mdistill text my friends are cool but they eat too many carbs inputs tokenizertext maxlength1024 returntensorsjax encoderoutputs model encodeinputs decoderstarttokenid model config decoderstarttokenid decoderinputids jnp onesinputs inputids shape0 1 dtypei4 decoderstarttokenid outputs model decodedecoderinputids encoderoutputs logits outputs logits outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict encoderhiddenstates encoderoutputs0 if encoderattentionmask is none batchsize sequencelength encoderhiddenstates shape 2 encoderattentionmask jnp onesbatchsize sequencelength batchsize sequencelength decoderinputids shape if decoderattentionmask is none decoderattentionmask jnp onesbatchsize sequencelength if decoderpositionids is none if pastkeyvalues is not none raise valueerrormake sure to provide decoderpositionids when passing pastkeyvalues decoderpositionids jnp broadcastto jnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng inputs params params or self params if pastkeyvalues are passed then cache is already initialized a private flag initcache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxblenderbotattention module if pastkeyvalues inputscache pastkeyvalues mutable cache else mutable false def decoderforwardmodule decoderinputids decoderattentionmask decoderpositionids kwargs decodermodule module getdecodermodule outputs decodermodule decoderinputids decoderattentionmask decoderpositionids kwargs hiddenstates outputs0 if self config tiewordembeddings sharedembedding module model variablesparamssharedembedding lmlogits module lmhead applyparams kernel sharedembedding t hiddenstates else lmlogits module lmheadhiddenstates lmlogits module finallogitsbias return lmlogits outputs outputs self module apply inputs decoderinputidsjnp arraydecoderinputids dtypei4 decoderattentionmaskjnp arraydecoderattentionmask dtypei4 decoderpositionidsjnp arraydecoderpositionids dtypei4 encoderhiddenstatesencoderhiddenstates encoderattentionmaskjnp arrayencoderattentionmask dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs mutablemutable methoddecoderforward if pastkeyvalues is none lmlogits decoderoutputs outputs else lmlogits decoderoutputs past outputs if returndict outputs flaxcausallmoutputwithcrossattentions logitslmlogits hiddenstatesdecoderoutputs hiddenstates attentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions else outputs lmlogits decoderoutputs1 add updated cache to model output if pastkeyvalues is not none and returndict outputspastkeyvalues unfreezepastcache return outputs elif pastkeyvalues is not none and not returndict outputs outputs 1 unfreezepastcache outputs1 return outputs def prepareinputsforgeneration self decoderinputids maxlength attentionmask optionaljax array none decoderattentionmask optionaljax array none encoderoutputsnone kwargs initializing the cache batchsize seqlength decoderinputids shape pastkeyvalues self initcachebatchsize maxlength encoderoutputs note that usually one would have to put 0 s in the attentionmask for x inputids shape1 and x cachelength but since the decoder uses a causal mask those positions are masked anyways thus we can create a single static attentionmask here which is more efficient for compilation extendedattentionmask jnp onesbatchsize maxlength dtypei4 if decoderattentionmask is not none positionids decoderattentionmask cumsumaxis1 1 extendedattentionmask lax dynamicupdatesliceextendedattentionmask decoderattentionmask 0 0 else positionids jnp broadcasttojnp arangeseqlength dtypei4none batchsize seqlength return pastkeyvalues pastkeyvalues encoderoutputs encoderoutputs encoderattentionmask attentionmask decoderattentionmask extendedattentionmask decoderpositionids positionids def updateinputsforgenerationself modeloutputs modelkwargs modelkwargspastkeyvalues modeloutputs pastkeyvalues modelkwargsdecoderpositionids modelkwargsdecoderpositionids 1 1 return modelkwargs flaxblenderbotconditionalgenerationdocstring r returns conversation example py from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration frompretrainedfacebookblenderbot400mdistill tokenizer autotokenizer frompretrainedfacebookblenderbot400mdistill utterance my friends are cool but they eat too many carbs inputs tokenizerutterance maxlength1024 returntensorsnp generate reply replyids model generateinputsinputids numbeams4 maxlength5 earlystoppingtrue sequences printtokenizer decodeg skipspecialtokenstrue cleanuptokenizationspacesfalse for g in replyids coding utf 8 2021 the fairseq s and 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 flax blenderbot model this model inherits from flaxpretrainedmodel 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 flax linen flax nn module https flax readthedocs io en latest _autosummary flax nn module html subclass use it as a regular flax module and refer to the flax documentation for all matter related to general usage and behavior finally this model supports inherent jax features such as just in time jit compilation https jax readthedocs io en latest jax html just in time compilation jit automatic differentiation https jax readthedocs io en latest jax html automatic differentiation vectorization https jax readthedocs io en latest jax html vectorization vmap parallelization https jax readthedocs io en latest jax html parallelization pmap parameters config blenderbotconfig 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 flaxpretrainedmodel from_pretrained method to load the model weights args input_ids jnp ndarray of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask jnp ndarray of shape batch_size sequence_length 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 attention mask decoder_input_ids jnp ndarray of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids for translation and summarization training decoder_input_ids should be provided if no decoder_input_ids is provided the model will create this tensor by shifting the input_ids to the right for denoising pre training following the paper decoder_attention_mask jnp ndarray of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default if you want to change padding behavior you should modify to your needs see diagram 1 in the paper https arxiv org abs 1910 13461 for more information on the default strategy position_ids numpy ndarray of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 decoder_position_ids numpy ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 output_attentions 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 bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple args input_ids jnp ndarray of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask jnp ndarray of shape batch_size sequence_length 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 attention mask position_ids numpy ndarray of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 output_attentions 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 bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple args decoder_input_ids jnp ndarray of shape batch_size target_sequence_length indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids for translation and summarization training decoder_input_ids should be provided if no decoder_input_ids is provided the model will create this tensor by shifting the input_ids to the right for denoising pre training following the paper encoder_outputs tuple tuple jnp ndarray tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask jnp ndarray of shape batch_size sequence_length 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 attention mask decoder_attention_mask jnp ndarray of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default if you want to change padding behavior you should modify to your needs see diagram 1 in the paper https arxiv org abs 1910 13461 for more information on the default strategy decoder_position_ids numpy ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config max_position_embeddings 1 past_key_values dict str np ndarray optional returned by init_cache or when passing previous past_key_values dictionary of pre computed hidden states key and values in the attention blocks that can be used for fast auto regressive decoding pre computed key and value hidden states are of shape batch_size max_length output_attentions 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 bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models bart modeling_flax_bart shift_tokens_right shift input ids one token to the right copied from transformers models bart modeling_flax_bart flaxbartattention with bart blenderbot the dtype of the computation this function takes projected key value states from a single input token and concatenates the states to cached states from previous steps this function is slighly adapted from the official flax repository https github com google flax blob 491ce18759622506588784b4fca0e4bf05f8c8cd flax linen attention py l252 detect if we re initializing by absence of existing cache data update key value caches with our new 1d spatial slices causal mask for cached decoder self attention our single query position should only attend to those key positions that have already been generated and cached not the remaining zero elements input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder get query proj get key value proj cross_attentions self_attention handle cache prepare causal attention mask combine masks if needed during fast autoregressive decoding we feed one position at a time and cache the keys and values step by step convert the boolean attention mask to an attention bias attention mask in the form of attention bias copied from transformers models mbart modeling_flax_mbart flaxmbartencoderlayer with mbart blenderbot copied from transformers models bart modeling_flax_bart flaxbartencoderlayercollection with bart blenderbot the dtype of the computation add layerdrop see https arxiv org abs 1909 11556 for description skip the layer copied from transformers models mbart modeling_flax_mbart flaxmbartdecoderlayer with mbart blenderbot self attention cross attention block fully connected copied from transformers models bart modeling_flax_bart flaxbartdecoderlayercollection with bart blenderbot the dtype of the computation decoder layers add layerdrop see https arxiv org abs 1909 11556 for description add hidden states from the last decoder layer the dtype of the computation update the last element in hidden_states after applying layernorm above the dtype of the computation embed positions update the last element in hidden_states after applying layernorm above copied from transformers models bart modeling_flax_bart flaxbartmodule with bart blenderbot the dtype of the computation init input tensors make sure initialization pass will work for flaxblenderbotforsequenceclassificationmodule args batch_size int batch_size used for fast auto regressive decoding defines the batch size of the initialized cache max_length int maximum possible length for auto regressive decoding defines the sequence length of the initialized cache encoder_outputs union flaxbasemodeloutput tuple tuple jnp ndarray encoder_outputs consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder init input variables to retrieve cache we only need to call the decoder to init the cache returns example python from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration from_pretrained facebook blenderbot 400m distill tokenizer autotokenizer from_pretrained facebook blenderbot 400m distill text my friends are cool but they eat too many carbs inputs tokenizer text max_length 1024 return_tensors jax encoder_outputs model encode inputs handle any prng if needed returns example python import jax numpy as jnp from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration from_pretrained facebook blenderbot 400m distill tokenizer autotokenizer from_pretrained facebook blenderbot 400m distill text my friends are cool but they eat too many carbs inputs tokenizer text max_length 1024 return_tensors jax encoder_outputs model encode inputs decoder_start_token_id model config decoder_start_token_id decoder_input_ids jnp ones inputs input_ids shape 0 1 dtype i4 decoder_start_token_id outputs model decode decoder_input_ids encoder_outputs last_decoder_hidden_states outputs last_hidden_state handle any prng if needed if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxblenderbotattention module add updated cache to model output prepare encoder inputs prepare decoder inputs handle any prng if needed the dtype of the computation copied from transformers models bart modeling_flax_bart flaxbartforconditionalgenerationmodule with bart blenderbot returns example python import jax numpy as jnp from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration from_pretrained facebook blenderbot 400m distill tokenizer autotokenizer from_pretrained facebook blenderbot 400m distill text my friends are cool but they eat too many carbs inputs tokenizer text max_length 1024 return_tensors jax encoder_outputs model encode inputs decoder_start_token_id model config decoder_start_token_id decoder_input_ids jnp ones inputs input_ids shape 0 1 dtype i4 decoder_start_token_id outputs model decode decoder_input_ids encoder_outputs logits outputs logits handle any prng if needed if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxblenderbotattention module add updated cache to model output initializing the cache note that usually one would have to put 0 s in the attention_mask for x input_ids shape 1 and x cache_length but since the decoder uses a causal mask those positions are masked anyways thus we can create a single static attention_mask here which is more efficient for compilation returns conversation example py from transformers import autotokenizer flaxblenderbotforconditionalgeneration model flaxblenderbotforconditionalgeneration from_pretrained facebook blenderbot 400m distill tokenizer autotokenizer from_pretrained facebook blenderbot 400m distill utterance my friends are cool but they eat too many carbs inputs tokenizer utterance max_length 1024 return_tensors np generate reply reply_ids model generate inputs input_ids num_beams 4 max_length 5 early_stopping true sequences print tokenizer decode g skip_special_tokens true clean_up_tokenization_spaces false for g in reply_ids
import math
import random
from functools import partial
from typing import Callable, Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
from flax.linen import combine_masks, make_causal_mask
from flax.linen.attention import dot_product_attention_weights
from flax.traverse_util import flatten_dict, unflatten_dict
from jax import lax
from jax.random import PRNGKey

from ...modeling_flax_outputs import (
    FlaxBaseModelOutput,
    FlaxBaseModelOutputWithPastAndCrossAttentions,
    FlaxCausalLMOutputWithCrossAttentions,
    FlaxSeq2SeqLMOutput,
    FlaxSeq2SeqModelOutput,
)
from ...modeling_flax_utils import (
    ACT2FN,
    FlaxPreTrainedModel,
    append_call_sample_docstring,
    append_replace_return_docstrings,
    overwrite_call_docstring,
)
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_blenderbot import BlenderbotConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "BlenderbotConfig"
_CHECKPOINT_FOR_DOC = "facebook/blenderbot-400M-distill"


BLENDERBOT_START_DOCSTRING = r

BLENDERBOT_INPUTS_DOCSTRING = r


BLENDERBOT_ENCODE_INPUTS_DOCSTRING = r

BLENDERBOT_DECODE_INPUTS_DOCSTRING = r



def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:
    
    shifted_input_ids = jnp.zeros_like(input_ids)
    shifted_input_ids = shifted_input_ids.at[:, 1:].set(input_ids[:, :-1])
    shifted_input_ids = shifted_input_ids.at[:, 0].set(decoder_start_token_id)

    shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
    return shifted_input_ids



class FlaxBlenderbotAttention(nn.Module):
    config: BlenderbotConfig
    embed_dim: int
    num_heads: int
    dropout: float = 0.0
    causal: bool = False
    bias: bool = True
    dtype: jnp.dtype = jnp.float32  

    def setup(self) -> None:
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {self.num_heads})."
            )

        dense = partial(
            nn.Dense,
            self.embed_dim,
            use_bias=self.bias,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std),
        )

        self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()
        self.out_proj = dense()

        self.dropout_layer = nn.Dropout(rate=self.dropout)

        if self.causal:
            self.causal_mask = make_causal_mask(
                jnp.ones((1, self.config.max_position_embeddings), dtype="bool"), dtype="bool"
            )

    def _split_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))

    def _merge_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))

    @nn.compact
    def _concatenate_to_cache(self, key, value, query, attention_mask):
        
        
        is_initialized = self.has_variable("cache", "cached_key")
        cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
        cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
        cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))

        if is_initialized:
            *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
            
            cur_index = cache_index.value
            indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
            key = lax.dynamic_update_slice(cached_key.value, key, indices)
            value = lax.dynamic_update_slice(cached_value.value, value, indices)
            cached_key.value = key
            cached_value.value = value
            num_updated_cache_vectors = query.shape[1]
            cache_index.value = cache_index.value + num_updated_cache_vectors
            
            pad_mask = jnp.broadcast_to(
                jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
                tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
            )
            attention_mask = combine_masks(pad_mask, attention_mask)
        return key, value, attention_mask

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        key_value_states: Optional[jnp.ndarray] = None,
        attention_mask: Optional[jnp.ndarray] = None,
        init_cache: bool = False,
        deterministic: bool = True,
    ) -> Tuple[jnp.ndarray]:
        

        
        
        is_cross_attention = key_value_states is not None
        batch_size = hidden_states.shape[0]

        
        query_states = self.q_proj(hidden_states)
        
        if is_cross_attention:
            
            key_states = self.k_proj(key_value_states)
            value_states = self.v_proj(key_value_states)
        else:
            
            key_states = self.k_proj(hidden_states)
            value_states = self.v_proj(hidden_states)

        query_states = self._split_heads(query_states)
        key_states = self._split_heads(key_states)
        value_states = self._split_heads(value_states)

        
        if self.causal:
            query_length, key_length = query_states.shape[1], key_states.shape[1]
            if self.has_variable("cache", "cached_key"):
                mask_shift = self.variables["cache"]["cache_index"]
                max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
                causal_mask = lax.dynamic_slice(
                    self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
                )
            else:
                causal_mask = self.causal_mask[:, :, :query_length, :key_length]
            causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])

        
        if attention_mask is not None and self.causal:
            attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
            attention_mask = combine_masks(attention_mask, causal_mask)
        elif self.causal:
            attention_mask = causal_mask
        elif attention_mask is not None:
            attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))

        
        
        if self.causal and (self.has_variable("cache", "cached_key") or init_cache):
            key_states, value_states, attention_mask = self._concatenate_to_cache(
                key_states, value_states, query_states, attention_mask
            )

        
        if attention_mask is not None:
            
            attention_bias = lax.select(
                attention_mask > 0,
                jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
                jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
            )
        else:
            attention_bias = None

        dropout_rng = None
        if not deterministic and self.dropout > 0.0:
            dropout_rng = self.make_rng("dropout")

        attn_weights = dot_product_attention_weights(
            query_states,
            key_states,
            bias=attention_bias,
            dropout_rng=dropout_rng,
            dropout_rate=self.dropout,
            broadcast_dropout=True,
            deterministic=deterministic,
            dtype=self.dtype,
            precision=None,
        )

        attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)
        attn_output = self._merge_heads(attn_output)
        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights



class FlaxBlenderbotEncoderLayer(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.embed_dim = self.config.d_model
        self.self_attn = FlaxBlenderbotAttention(
            config=self.config,
            embed_dim=self.embed_dim,
            num_heads=self.config.encoder_attention_heads,
            dropout=self.config.attention_dropout,
            dtype=self.dtype,
        )
        self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)
        self.activation_fn = ACT2FN[self.config.activation_function]
        self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)
        self.fc1 = nn.Dense(
            self.config.encoder_ffn_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std),
        )
        self.fc2 = nn.Dense(
            self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)
        )
        self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        attention_mask: jnp.ndarray,
        output_attentions: bool = True,
        deterministic: bool = True,
    ) -> Tuple[jnp.ndarray]:
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask)
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs



class FlaxBlenderbotEncoderLayerCollection(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32  

    def setup(self):
        self.layers = [
            FlaxBlenderbotEncoderLayer(self.config, name=str(i), dtype=self.dtype)
            for i in range(self.config.encoder_layers)
        ]
        self.layerdrop = self.config.encoder_layerdrop

    def __call__(
        self,
        hidden_states,
        attention_mask,
        deterministic: bool = True,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        all_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None

        for encoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            
            dropout_probability = random.uniform(0, 1)
            if not deterministic and (dropout_probability < self.layerdrop):  
                layer_outputs = (None, None)
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    output_attentions,
                    deterministic,
                )
            hidden_states = layer_outputs[0]
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        outputs = (hidden_states, all_hidden_states, all_attentions)

        if not return_dict:
            return tuple(v for v in outputs if v is not None)

        return FlaxBaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
        )



class FlaxBlenderbotDecoderLayer(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.embed_dim = self.config.d_model
        self.self_attn = FlaxBlenderbotAttention(
            config=self.config,
            embed_dim=self.embed_dim,
            num_heads=self.config.decoder_attention_heads,
            dropout=self.config.attention_dropout,
            causal=True,
            dtype=self.dtype,
        )
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)
        self.activation_fn = ACT2FN[self.config.activation_function]
        self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)

        self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)
        self.encoder_attn = FlaxBlenderbotAttention(
            config=self.config,
            embed_dim=self.embed_dim,
            num_heads=self.config.decoder_attention_heads,
            dropout=self.config.attention_dropout,
            dtype=self.dtype,
        )
        self.encoder_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)
        self.fc1 = nn.Dense(
            self.config.decoder_ffn_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std),
        )
        self.fc2 = nn.Dense(
            self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)
        )
        self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)

    def __call__(
        self,
        hidden_states: jnp.ndarray,
        attention_mask: jnp.ndarray,
        encoder_hidden_states: Optional[jnp.ndarray] = None,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        init_cache: bool = False,
        output_attentions: bool = True,
        deterministic: bool = True,
    ) -> Tuple[jnp.ndarray]:
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states, attention_mask=attention_mask, init_cache=init_cache
        )
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states

            hidden_states = self.encoder_attn_layer_norm(hidden_states)
            hidden_states, cross_attn_weights = self.encoder_attn(
                hidden_states=hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
            )
            hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
            hidden_states = residual + hidden_states

        
        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights, cross_attn_weights)

        return outputs



class FlaxBlenderbotDecoderLayerCollection(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32  

    def setup(self):
        self.layers = [
            FlaxBlenderbotDecoderLayer(self.config, name=str(i), dtype=self.dtype)
            for i in range(self.config.decoder_layers)
        ]
        self.layerdrop = self.config.decoder_layerdrop

    def __call__(
        self,
        hidden_states,
        attention_mask,
        encoder_hidden_states: Optional[jnp.ndarray] = None,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        deterministic: bool = True,
        init_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ):
        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None

        for decoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
                
            dropout_probability = random.uniform(0, 1)
            if not deterministic and (dropout_probability < self.layerdrop):
                layer_outputs = (None, None, None)
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    init_cache=init_cache,
                    output_attentions=output_attentions,
                    deterministic=deterministic,
                )

            hidden_states = layer_outputs[0]
            if output_attentions:
                all_self_attns += (layer_outputs[1],)

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        outputs = [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions]

        if not return_dict:
            return tuple(v for v in outputs if v is not None)

        return FlaxBaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )


class FlaxBlenderbotEncoder(nn.Module):
    config: BlenderbotConfig
    embed_tokens: nn.Embed
    dtype: jnp.dtype = jnp.float32  

    def setup(self):
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)

        embed_dim = self.config.d_model
        self.padding_idx = self.config.pad_token_id
        self.max_source_positions = self.config.max_position_embeddings
        self.embed_scale = math.sqrt(embed_dim) if self.config.scale_embedding else 1.0

        self.embed_positions = nn.Embed(
            self.config.max_position_embeddings,
            embed_dim,
            embedding_init=jax.nn.initializers.normal(self.config.init_std),
        )
        self.layers = FlaxBlenderbotEncoderLayerCollection(self.config, self.dtype)
        self.layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)

    def __call__(
        self,
        input_ids,
        attention_mask,
        position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        input_shape = input_ids.shape
        input_ids = input_ids.reshape(-1, input_shape[-1])

        inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale

        embed_pos = self.embed_positions(position_ids)

        hidden_states = inputs_embeds + embed_pos
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)

        outputs = self.layers(
            hidden_states,
            attention_mask,
            deterministic=deterministic,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        last_hidden_states = outputs[0]
        last_hidden_states = self.layer_norm(last_hidden_states)

        
        hidden_states = None
        if output_hidden_states:
            hidden_states = outputs[1]
            hidden_states = hidden_states[:-1] + (last_hidden_states,)

        if not return_dict:
            outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
            return tuple(v for v in outputs if v is not None)

        return FlaxBaseModelOutput(
            last_hidden_state=last_hidden_states,
            hidden_states=hidden_states,
            attentions=outputs.attentions,
        )


class FlaxBlenderbotDecoder(nn.Module):
    config: BlenderbotConfig
    embed_tokens: nn.Embed
    dtype: jnp.dtype = jnp.float32  

    def setup(self):
        self.dropout_layer = nn.Dropout(rate=self.config.dropout)

        embed_dim = self.config.d_model
        self.padding_idx = self.config.pad_token_id
        self.max_target_positions = self.config.max_position_embeddings
        self.embed_scale = math.sqrt(self.config.d_model) if self.config.scale_embedding else 1.0

        self.embed_positions = nn.Embed(
            self.config.max_position_embeddings,
            embed_dim,
            embedding_init=jax.nn.initializers.normal(self.config.init_std),
        )

        self.layers = FlaxBlenderbotDecoderLayerCollection(self.config, self.dtype)
        self.layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)

    def __call__(
        self,
        input_ids,
        attention_mask,
        position_ids,
        encoder_hidden_states: Optional[jnp.ndarray] = None,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        init_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        input_shape = input_ids.shape
        input_ids = input_ids.reshape(-1, input_shape[-1])

        inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale

        
        positions = self.embed_positions(position_ids)

        hidden_states = inputs_embeds + positions
        hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)

        outputs = self.layers(
            hidden_states,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            deterministic=deterministic,
            init_cache=init_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        last_hidden_states = outputs[0]
        last_hidden_states = self.layer_norm(last_hidden_states)

        
        hidden_states = None
        if output_hidden_states:
            hidden_states = outputs[1]
            hidden_states = hidden_states[:-1] + (last_hidden_states,)

        if not return_dict:
            outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
            return tuple(v for v in outputs if v is not None)

        return FlaxBaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=last_hidden_states,
            hidden_states=hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )



class FlaxBlenderbotModule(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32  

    def setup(self):
        self.shared = nn.Embed(
            self.config.vocab_size,
            self.config.d_model,
            embedding_init=jax.nn.initializers.normal(self.config.init_std),
            dtype=self.dtype,
        )

        self.encoder = FlaxBlenderbotEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)
        self.decoder = FlaxBlenderbotDecoder(self.config, dtype=self.dtype, embed_tokens=self.shared)

    def _get_encoder_module(self):
        return self.encoder

    def _get_decoder_module(self):
        return self.decoder

    def __call__(
        self,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        position_ids,
        decoder_position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        encoder_outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            position_ids=decoder_position_ids,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return FlaxSeq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )


class FlaxBlenderbotPreTrainedModel(FlaxPreTrainedModel):
    config_class = BlenderbotConfig
    base_model_prefix: str = "model"
    module_class: nn.Module = None

    def __init__(
        self,
        config: BlenderbotConfig,
        input_shape: Tuple[int] = (1, 1),
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
        **kwargs,
    ):
        module = self.module_class(config=config, dtype=dtype, **kwargs)
        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
        
        input_ids = jnp.zeros(input_shape, dtype="i4")
        
        input_ids = input_ids.at[(..., -1)].set(self.config.eos_token_id)
        attention_mask = jnp.ones_like(input_ids)
        decoder_input_ids = input_ids
        decoder_attention_mask = jnp.ones_like(input_ids)

        batch_size, sequence_length = input_ids.shape
        position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
        decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))

        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        random_params = self.module.init(
            rngs,
            input_ids,
            attention_mask,
            decoder_input_ids,
            decoder_attention_mask,
            position_ids,
            decoder_position_ids,
        )["params"]

        if params is not None:
            random_params = flatten_dict(unfreeze(random_params))
            params = flatten_dict(unfreeze(params))
            for missing_key in self._missing_keys:
                params[missing_key] = random_params[missing_key]
            self._missing_keys = set()
            return freeze(unflatten_dict(params))
        else:
            return random_params

    def init_cache(self, batch_size, max_length, encoder_outputs):
        r
        
        decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
        decoder_attention_mask = jnp.ones_like(decoder_input_ids)
        decoder_position_ids = jnp.broadcast_to(
            jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape
        )

        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
            decoder_module = module._get_decoder_module()
            return decoder_module(
                decoder_input_ids,
                decoder_attention_mask,
                decoder_position_ids,
                **kwargs,
            )

        init_variables = self.module.init(
            jax.random.PRNGKey(0),
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            decoder_position_ids=decoder_position_ids,
            encoder_hidden_states=encoder_outputs[0],
            init_cache=True,
            method=_decoder_forward,  
        )
        return unfreeze(init_variables["cache"])

    @add_start_docstrings(BLENDERBOT_ENCODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=BlenderbotConfig)
    def encode(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        position_ids: Optional[jnp.ndarray] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)
        if position_ids is None:
            batch_size, sequence_length = input_ids.shape
            position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))

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

        def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):
            encode_module = module._get_encoder_module()
            return encode_module(input_ids, attention_mask, position_ids, **kwargs)

        return self.module.apply(
            {"params": params or self.params},
            input_ids=jnp.array(input_ids, dtype="i4"),
            attention_mask=jnp.array(attention_mask, dtype="i4"),
            position_ids=jnp.array(position_ids, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
            method=_encoder_forward,
        )

    @add_start_docstrings(BLENDERBOT_DECODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(
        output_type=FlaxBaseModelOutputWithPastAndCrossAttentions, config_class=BlenderbotConfig
    )
    def decode(
        self,
        decoder_input_ids,
        encoder_outputs,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        past_key_values: dict = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        encoder_hidden_states = encoder_outputs[0]
        if encoder_attention_mask is None:
            batch_size, sequence_length = encoder_hidden_states.shape[:2]
            encoder_attention_mask = jnp.ones((batch_size, sequence_length))

        batch_size, sequence_length = decoder_input_ids.shape
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones((batch_size, sequence_length))

        if decoder_position_ids is None:
            if past_key_values is not None:
                raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")

            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

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

        inputs = {"params": params or self.params}

        
        
        
        if past_key_values:
            inputs["cache"] = past_key_values
            mutable = ["cache"]
        else:
            mutable = False

        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
            decoder_module = module._get_decoder_module()
            return decoder_module(
                decoder_input_ids,
                decoder_attention_mask,
                decoder_position_ids,
                **kwargs,
            )

        outputs = self.module.apply(
            inputs,
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
            mutable=mutable,
            method=_decoder_forward,
        )

        
        if past_key_values is not None and return_dict:
            outputs, past = outputs
            outputs["past_key_values"] = unfreeze(past["cache"])
            return outputs
        elif past_key_values is not None and not return_dict:
            outputs, past = outputs
            outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]

        return outputs

    @add_start_docstrings_to_model_forward(BLENDERBOT_INPUTS_DOCSTRING)
    def __call__(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        decoder_input_ids: Optional[jnp.ndarray] = None,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        position_ids: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        
        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)
        if position_ids is None:
            batch_size, sequence_length = input_ids.shape
            position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))

        
        if decoder_input_ids is None:
            decoder_input_ids = shift_tokens_right(
                input_ids, self.config.pad_token_id, decoder_start_token_id=self.config.decoder_start_token_id
            )
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones_like(decoder_input_ids)
        if decoder_position_ids is None:
            batch_size, sequence_length = decoder_input_ids.shape
            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

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

        return self.module.apply(
            {"params": params or self.params},
            input_ids=jnp.array(input_ids, dtype="i4"),
            attention_mask=jnp.array(attention_mask, dtype="i4"),
            position_ids=jnp.array(position_ids, dtype="i4"),
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
        )


@add_start_docstrings(
    "The bare MBart Model transformer outputting raw hidden-states without any specific head on top.",
    BLENDERBOT_START_DOCSTRING,
)
class FlaxBlenderbotModel(FlaxBlenderbotPreTrainedModel):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32  
    module_class = FlaxBlenderbotModule


append_call_sample_docstring(FlaxBlenderbotModel, _CHECKPOINT_FOR_DOC, FlaxSeq2SeqModelOutput, _CONFIG_FOR_DOC)



class FlaxBlenderbotForConditionalGenerationModule(nn.Module):
    config: BlenderbotConfig
    dtype: jnp.dtype = jnp.float32
    bias_init: Callable[..., jnp.ndarray] = jax.nn.initializers.zeros

    def setup(self):
        self.model = FlaxBlenderbotModule(config=self.config, dtype=self.dtype)
        self.lm_head = nn.Dense(
            self.model.shared.num_embeddings,
            use_bias=False,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std),
        )
        self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, self.model.shared.num_embeddings))

    def _get_encoder_module(self):
        return self.model.encoder

    def _get_decoder_module(self):
        return self.model.decoder

    def __call__(
        self,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        position_ids,
        decoder_position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            position_ids=position_ids,
            decoder_position_ids=decoder_position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        hidden_states = outputs[0]

        if self.config.tie_word_embeddings:
            shared_embedding = self.model.variables["params"]["shared"]["embedding"]
            lm_logits = self.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
        else:
            lm_logits = self.lm_head(hidden_states)

        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return output

        return FlaxSeq2SeqLMOutput(
            logits=lm_logits,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )


@add_start_docstrings(
    "The Blenderbot Model with a language modeling head. Can be used for summarization.", BLENDERBOT_START_DOCSTRING
)
class FlaxBlenderbotForConditionalGeneration(FlaxBlenderbotPreTrainedModel):
    module_class = FlaxBlenderbotForConditionalGenerationModule
    dtype: jnp.dtype = jnp.float32

    @add_start_docstrings(BLENDERBOT_DECODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=BlenderbotConfig)
    def decode(
        self,
        decoder_input_ids,
        encoder_outputs,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        past_key_values: dict = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        encoder_hidden_states = encoder_outputs[0]
        if encoder_attention_mask is None:
            batch_size, sequence_length = encoder_hidden_states.shape[:2]
            encoder_attention_mask = jnp.ones((batch_size, sequence_length))

        batch_size, sequence_length = decoder_input_ids.shape
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones((batch_size, sequence_length))

        if decoder_position_ids is None:
            if past_key_values is not None:
                raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")

            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

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

        inputs = {"params": params or self.params}

        
        
        
        if past_key_values:
            inputs["cache"] = past_key_values
            mutable = ["cache"]
        else:
            mutable = False

        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
            decoder_module = module._get_decoder_module()
            outputs = decoder_module(
                decoder_input_ids,
                decoder_attention_mask,
                decoder_position_ids,
                **kwargs,
            )
            hidden_states = outputs[0]

            if self.config.tie_word_embeddings:
                shared_embedding = module.model.variables["params"]["shared"]["embedding"]
                lm_logits = module.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
            else:
                lm_logits = module.lm_head(hidden_states)

            lm_logits += module.final_logits_bias
            return lm_logits, outputs

        outputs = self.module.apply(
            inputs,
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
            mutable=mutable,
            method=_decoder_forward,
        )

        if past_key_values is None:
            lm_logits, decoder_outputs = outputs
        else:
            (lm_logits, decoder_outputs), past = outputs

        if return_dict:
            outputs = FlaxCausalLMOutputWithCrossAttentions(
                logits=lm_logits,
                hidden_states=decoder_outputs.hidden_states,
                attentions=decoder_outputs.attentions,
                cross_attentions=decoder_outputs.cross_attentions,
            )
        else:
            outputs = (lm_logits,) + decoder_outputs[1:]

        
        if past_key_values is not None and return_dict:
            outputs["past_key_values"] = unfreeze(past["cache"])
            return outputs
        elif past_key_values is not None and not return_dict:
            outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]

        return outputs

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        max_length,
        attention_mask: Optional[jax.Array] = None,
        decoder_attention_mask: Optional[jax.Array] = None,
        encoder_outputs=None,
        **kwargs,
    ):
        
        batch_size, seq_length = decoder_input_ids.shape

        past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
        
        
        
        extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
        if decoder_attention_mask is not None:
            position_ids = decoder_attention_mask.cumsum(axis=-1) - 1
            extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
        else:
            position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))

        return {
            "past_key_values": past_key_values,
            "encoder_outputs": encoder_outputs,
            "encoder_attention_mask": attention_mask,
            "decoder_attention_mask": extended_attention_mask,
            "decoder_position_ids": position_ids,
        }

    def update_inputs_for_generation(self, model_outputs, model_kwargs):
        model_kwargs["past_key_values"] = model_outputs.past_key_values
        model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
        return model_kwargs


FLAX_BLENDERBOT_CONDITIONAL_GENERATION_DOCSTRING = r

overwrite_call_docstring(
    FlaxBlenderbotForConditionalGeneration,
    BLENDERBOT_INPUTS_DOCSTRING + FLAX_BLENDERBOT_CONDITIONAL_GENERATION_DOCSTRING,
)
append_replace_return_docstrings(
    FlaxBlenderbotForConditionalGeneration, output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC
)