import{S as wa,i as va,s as ba,e as a,k as c,w as m,t as n,M as ka,c as r,d as o,m as u,a as i,x as g,h as s,b as f,G as t,g as p,y as _,L as ya,q as w,o as v,B as b,v as Ta}from"../chunks/vendor-hf-doc-builder.js";import{I as vt}from"../chunks/IconCopyLink-hf-doc-builder.js";import{C}from"../chunks/CodeBlock-hf-doc-builder.js";import{D as Aa}from"../chunks/DocNotebookDropdown-hf-doc-builder.js";function $a(cs){let j,bt,L,M,Ge,K,go,Ie,_o,kt,J,yt,O,wo,Q,vo,bo,Tt,xe,ko,At,d,yo,Ne,To,Ao,Re,$o,xo,V,He,Eo,jo,Ye,Lo,Fo,Ue,Xo,qo,Be,zo,Po,Ke,Co,Mo,Je,Oo,Do,$t,k,Wo,Z,So,Go,ee,Io,No,te,Ro,Ho,oe,Yo,Uo,xt,Ee,Bo,Et,F,D,Qe,ne,Ko,Ve,Jo,jt,je,Qo,Lt,se,Ft,W,Vo,Ze,Zo,en,Xt,ae,qt,Le,tn,zt,re,Pt,T,on,et,nn,sn,tt,an,rn,Ct,ie,Mt,X,S,ot,le,ln,nt,pn,Ot,G,fn,st,cn,un,Dt,pe,Wt,Fe,dn,St,fe,Gt,I,hn,at,mn,gn,It,q,N,rt,ce,_n,it,wn,Nt,A,vn,lt,bn,kn,ue,yn,Tn,Rt,R,An,pt,$n,xn,Ht,$,En,ft,jn,Ln,ct,Fn,Xn,Yt,de,Ut,H,qn,ut,zn,Pn,Bt,he,Kt,Xe,Cn,Jt,me,Qt,Y,Mn,dt,On,Dn,Vt,qe,Wn,Zt,z,U,ht,ge,Sn,mt,Gn,eo,ze,In,to,y,x,_e,Nn,Rn,we,Hn,Yn,ve,Un,Bn,Kn,Pe,be,Jn,Qn,Vn,Ce,ke,Zn,es,ts,Me,os,P,gt,ye,ns,ss,_t,Te,as,rs,wt,Ae,is,oo;return K=new vt({}),J=new Aa({props:{classNames:"absolute z-10 right-0 top-0",options:[{label:"Mixed",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/transformers_doc/en/tf_xla.ipynb"},{label:"PyTorch",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/transformers_doc/en/pytorch/tf_xla.ipynb"},{label:"TensorFlow",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/transformers_doc/en/tensorflow/tf_xla.ipynb"},{label:"Mixed",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/transformers_doc/en/tf_xla.ipynb"},{label:"PyTorch",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/transformers_doc/en/pytorch/tf_xla.ipynb"},{label:"TensorFlow",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/transformers_doc/en/tensorflow/tf_xla.ipynb"}]}}),ne=new vt({}),se=new C({props:{code:`import tensorflow as tf

model = tf.keras.Sequential(
    [tf.keras.layers.Dense(10, input_shape=(10,), activation="relu"), tf.keras.layers.Dense(5, activation="softmax")]
)`,highlighted:`<span class="hljs-keyword">import</span> tensorflow <span class="hljs-keyword">as</span> tf

model = tf.keras.Sequential(
    [tf.keras.layers.Dense(<span class="hljs-number">10</span>, input_shape=(<span class="hljs-number">10</span>,), activation=<span class="hljs-string">&quot;relu&quot;</span>), tf.keras.layers.Dense(<span class="hljs-number">5</span>, activation=<span class="hljs-string">&quot;softmax&quot;</span>)]
)`}}),ae=new C({props:{code:`# Generate random inputs for the model.
batch_size = 16
input_vector_dim = 10
random_inputs = tf.random.normal((batch_size, input_vector_dim))

# Run a forward pass.
_ = model(random_inputs)`,highlighted:`<span class="hljs-comment"># Generate random inputs for the model.</span>
batch_size = <span class="hljs-number">16</span>
input_vector_dim = <span class="hljs-number">10</span>
random_inputs = tf.random.normal((batch_size, input_vector_dim))

<span class="hljs-comment"># Run a forward pass.</span>
_ = model(random_inputs)`}}),re=new C({props:{code:`xla_fn = tf.function(model, jit_compile=True)
_ = xla_fn(random_inputs)`,highlighted:`xla_fn = tf.function(model, jit_compile=<span class="hljs-literal">True</span>)
_ = xla_fn(random_inputs)`}}),ie=new C({props:{code:"my_xla_fn = tf.function(model.my_xla_fn, jit_compile=True)",highlighted:'my_xla_fn = tf.function(model.my_xla_fn, jit_compile=<span class="hljs-literal">True</span>)'}}),le=new vt({}),pe=new C({props:{code:"pip install transformers --upgrade",highlighted:"pip install transformers --upgrade"}}),fe=new C({props:{code:`import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# Will error if the minimal version of Transformers is not installed.
from transformers.utils import check_min_version

check_min_version("4.21.0")


tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
input_string = ["TensorFlow is"]

# One line to create an XLA generation function
xla_generate = tf.function(model.generate, jit_compile=True)

tokenized_input = tokenizer(input_string, return_tensors="tf")
generated_tokens = xla_generate(**tokenized_input, num_beams=2)

decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")
# Generated -- TensorFlow is an open-source, open-source, distributed-source application # framework for the`,highlighted:`<span class="hljs-keyword">import</span> tensorflow <span class="hljs-keyword">as</span> tf
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoTokenizer, TFAutoModelForCausalLM

<span class="hljs-comment"># Will error if the minimal version of Transformers is not installed.</span>
<span class="hljs-keyword">from</span> transformers.utils <span class="hljs-keyword">import</span> check_min_version

check_min_version(<span class="hljs-string">&quot;4.21.0&quot;</span>)


tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>, padding_side=<span class="hljs-string">&quot;left&quot;</span>, pad_token=<span class="hljs-string">&quot;&lt;/s&gt;&quot;</span>)
model = TFAutoModelForCausalLM.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>)
input_string = [<span class="hljs-string">&quot;TensorFlow is&quot;</span>]

<span class="hljs-comment"># One line to create an XLA generation function</span>
xla_generate = tf.function(model.generate, jit_compile=<span class="hljs-literal">True</span>)

tokenized_input = tokenizer(input_string, return_tensors=<span class="hljs-string">&quot;tf&quot;</span>)
generated_tokens = xla_generate(**tokenized_input, num_beams=<span class="hljs-number">2</span>)

decoded_text = tokenizer.decode(generated_tokens[<span class="hljs-number">0</span>], skip_special_tokens=<span class="hljs-literal">True</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;Generated -- <span class="hljs-subst">{decoded_text}</span>&quot;</span>)
<span class="hljs-comment"># Generated -- TensorFlow is an open-source, open-source, distributed-source application # framework for the</span>`}}),ce=new vt({}),de=new C({props:{code:`import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
input_string = ["TensorFlow is"]

xla_generate = tf.function(model.generate, jit_compile=True)

# Here, we call the tokenizer with padding options.
tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")

generated_tokens = xla_generate(**tokenized_input, num_beams=2)
decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")`,highlighted:`<span class="hljs-keyword">import</span> tensorflow <span class="hljs-keyword">as</span> tf
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>, padding_side=<span class="hljs-string">&quot;left&quot;</span>, pad_token=<span class="hljs-string">&quot;&lt;/s&gt;&quot;</span>)
model = TFAutoModelForCausalLM.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>)
input_string = [<span class="hljs-string">&quot;TensorFlow is&quot;</span>]

xla_generate = tf.function(model.generate, jit_compile=<span class="hljs-literal">True</span>)

<span class="hljs-comment"># Here, we call the tokenizer with padding options.</span>
tokenized_input = tokenizer(input_string, pad_to_multiple_of=<span class="hljs-number">8</span>, padding=<span class="hljs-literal">True</span>, return_tensors=<span class="hljs-string">&quot;tf&quot;</span>)

generated_tokens = xla_generate(**tokenized_input, num_beams=<span class="hljs-number">2</span>)
decoded_text = tokenizer.decode(generated_tokens[<span class="hljs-number">0</span>], skip_special_tokens=<span class="hljs-literal">True</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;Generated -- <span class="hljs-subst">{decoded_text}</span>&quot;</span>)`}}),he=new C({props:{code:`import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")

xla_generate = tf.function(model.generate, jit_compile=True)

for input_string in ["TensorFlow is", "TensorFlow is a", "TFLite is a"]:
    tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")
    start = time.time_ns()
    generated_tokens = xla_generate(**tokenized_input, num_beams=2)
    end = time.time_ns()
    print(f"Execution time -- {(end - start) / 1e6:.1f} ms\\n")`,highlighted:`<span class="hljs-keyword">import</span> time
<span class="hljs-keyword">import</span> tensorflow <span class="hljs-keyword">as</span> tf
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>, padding_side=<span class="hljs-string">&quot;left&quot;</span>, pad_token=<span class="hljs-string">&quot;&lt;/s&gt;&quot;</span>)
model = TFAutoModelForCausalLM.from_pretrained(<span class="hljs-string">&quot;gpt2&quot;</span>)

xla_generate = tf.function(model.generate, jit_compile=<span class="hljs-literal">True</span>)

<span class="hljs-keyword">for</span> input_string <span class="hljs-keyword">in</span> [<span class="hljs-string">&quot;TensorFlow is&quot;</span>, <span class="hljs-string">&quot;TensorFlow is a&quot;</span>, <span class="hljs-string">&quot;TFLite is a&quot;</span>]:
    tokenized_input = tokenizer(input_string, pad_to_multiple_of=<span class="hljs-number">8</span>, padding=<span class="hljs-literal">True</span>, return_tensors=<span class="hljs-string">&quot;tf&quot;</span>)
    start = time.time_ns()
    generated_tokens = xla_generate(**tokenized_input, num_beams=<span class="hljs-number">2</span>)
    end = time.time_ns()
    <span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;Execution time -- <span class="hljs-subst">{(end - start) / <span class="hljs-number">1e6</span>:<span class="hljs-number">.1</span>f}</span> ms\\n&quot;</span>)`}}),me=new C({props:{code:`Execution time -- 30819.6 ms

Execution time -- 79.0 ms

Execution time -- 78.9 ms`,highlighted:`Execution time -- 30819.6 ms

Execution time -- 79.0 ms

Execution time -- 78.9 ms`}}),ge=new vt({}),{c(){j=a("meta"),bt=c(),L=a("h1"),M=a("a"),Ge=a("span"),m(K.$$.fragment),go=c(),Ie=a("span"),_o=n("XLA Integration for TensorFlow Models"),kt=c(),m(J.$$.fragment),yt=c(),O=a("p"),wo=n("Accelerated Linear Algebra, dubbed XLA, is a compiler for accelerating the runtime of TensorFlow Models. From the "),Q=a("a"),vo=n("official documentation"),bo=n(":"),Tt=c(),xe=a("p"),ko=n("XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes."),At=c(),d=a("p"),yo=n("Using XLA in TensorFlow is simple \u2013 it comes packaged inside the "),Ne=a("code"),To=n("tensorflow"),Ao=n(" library, and it can be triggered with the "),Re=a("code"),$o=n("jit_compile"),xo=n(" argument in any graph-creating function such as "),V=a("a"),He=a("code"),Eo=n("tf.function"),jo=n(". When using Keras methods like "),Ye=a("code"),Lo=n("fit()"),Fo=n(" and "),Ue=a("code"),Xo=n("predict()"),qo=n(", you can enable XLA simply by passing the "),Be=a("code"),zo=n("jit_compile"),Po=n(" argument to "),Ke=a("code"),Co=n("model.compile()"),Mo=n(". However, XLA is not limited to these methods - it can also be used to accelerate any arbitrary "),Je=a("code"),Oo=n("tf.function"),Do=n("."),$t=c(),k=a("p"),Wo=n("Several TensorFlow methods in \u{1F917} Transformers have been rewritten to be XLA-compatible, including text generation for models such as "),Z=a("a"),So=n("GPT2"),Go=n(", "),ee=a("a"),Io=n("T5"),No=n(" and "),te=a("a"),Ro=n("OPT"),Ho=n(", as well as speech processing for models such as "),oe=a("a"),Yo=n("Whisper"),Uo=n("."),xt=c(),Ee=a("p"),Bo=n("While the exact amount of speed-up is very much model-dependent, for TensorFlow text generation models inside \u{1F917} Transformers, we noticed a speed-up of ~100x. This document will explain how you can use XLA for these models to get the maximum amount of performance. We\u2019ll also provide links to additional resources if you\u2019re interested to learn more about the benchmarks and our design philosophy behind the XLA integration."),Et=c(),F=a("h2"),D=a("a"),Qe=a("span"),m(ne.$$.fragment),Ko=c(),Ve=a("span"),Jo=n("Running TF functions with XLA"),jt=c(),je=a("p"),Qo=n("Let us consider the following model in TensorFlow:"),Lt=c(),m(se.$$.fragment),Ft=c(),W=a("p"),Vo=n("The above model accepts inputs having a dimension of "),Ze=a("code"),Zo=n("(10, )"),en=n(". We can use the model for running a forward pass like so:"),Xt=c(),m(ae.$$.fragment),qt=c(),Le=a("p"),tn=n("In order to run the forward pass with an XLA-compiled function, we\u2019d need to do:"),zt=c(),m(re.$$.fragment),Pt=c(),T=a("p"),on=n("The default "),et=a("code"),nn=n("call()"),sn=n(" function of the "),tt=a("code"),an=n("model"),rn=n(" is used for compiling the XLA graph. But if there\u2019s any other model function you want to compile into XLA that\u2019s also possible with:"),Ct=c(),m(ie.$$.fragment),Mt=c(),X=a("h2"),S=a("a"),ot=a("span"),m(le.$$.fragment),ln=c(),nt=a("span"),pn=n("Running a TF text generation model with XLA from \u{1F917} Transformers"),Ot=c(),G=a("p"),fn=n("To enable XLA-accelerated generation within \u{1F917} Transformers, you need to have a recent version of "),st=a("code"),cn=n("transformers"),un=n(" installed. You can install it by running:"),Dt=c(),m(pe.$$.fragment),Wt=c(),Fe=a("p"),dn=n("And then you can run the following code:"),St=c(),m(fe.$$.fragment),Gt=c(),I=a("p"),hn=n("As you can notice, enabling XLA on "),at=a("code"),mn=n("generate()"),gn=n(" is just a single line of code. The rest of the code remains unchanged. However, there are a couple of gotchas in the above code snippet that are specific to XLA. You need to be aware of those to realize the speed-ups that XLA can bring in. We discuss these in the following section."),It=c(),q=a("h2"),N=a("a"),rt=a("span"),m(ce.$$.fragment),_n=c(),it=a("span"),wn=n("Gotchas to be aware of"),Nt=c(),A=a("p"),vn=n("When you are executing an XLA-enabled function (like "),lt=a("code"),bn=n("xla_generate()"),kn=n(" above) for the first time, it will internally try to infer the computation graph, which is time-consuming.  This process is known as "),ue=a("a"),yn=n("\u201Ctracing\u201D"),Tn=n("."),Rt=c(),R=a("p"),An=n("You might notice that the generation time is not fast. Successive calls of "),pt=a("code"),$n=n("xla_generate()"),xn=n(" (or any other XLA-enabled function) won\u2019t have to infer the computation graph, given the inputs to the function follow the same shape with which the computation graph was initially built. While this is not a problem for modalities with fixed input shapes (e.g., images), you must pay attention if you are working with variable input shape modalities (e.g., text)."),Ht=c(),$=a("p"),En=n("To ensure "),ft=a("code"),jn=n("xla_generate()"),Ln=n(" always operates with the same input shapes, you can specify the "),ct=a("code"),Fn=n("padding"),Xn=n(" arguments when calling the tokenizer."),Yt=c(),m(de.$$.fragment),Ut=c(),H=a("p"),qn=n("This way, you can ensure that the inputs to "),ut=a("code"),zn=n("xla_generate()"),Pn=n(" will always receive inputs with the shape it was traced with and thus leading to speed-ups in the generation time. You can verify this with the code below:"),Bt=c(),m(he.$$.fragment),Kt=c(),Xe=a("p"),Cn=n("On a Tesla T4 GPU, you can expect the outputs like so:"),Jt=c(),m(me.$$.fragment),Qt=c(),Y=a("p"),Mn=n("The first call to "),dt=a("code"),On=n("xla_generate()"),Dn=n(" is time-consuming because of tracing, but the successive calls are orders of magnitude faster. Keep in mind that any change in the generation options at any point with trigger re-tracing and thus leading to slow-downs in the generation time."),Vt=c(),qe=a("p"),Wn=n("We didn\u2019t cover all the text generation options \u{1F917} Transformers provides in this document. We encourage you to read the documentation for advanced use cases."),Zt=c(),z=a("h2"),U=a("a"),ht=a("span"),m(ge.$$.fragment),Sn=c(),mt=a("span"),Gn=n("Additional Resources"),eo=c(),ze=a("p"),In=n("Here, we leave you with some additional resources if you want to delve deeper into XLA in \u{1F917} Transformers and in general."),to=c(),y=a("ul"),x=a("li"),_e=a("a"),Nn=n("This Colab Notebook"),Rn=n(" provides an interactive demonstration if you want to fiddle with the XLA-compatible encoder-decoder (like "),we=a("a"),Hn=n("T5"),Yn=n(") and decoder-only (like "),ve=a("a"),Un=n("GPT2"),Bn=n(") text generation models."),Kn=c(),Pe=a("li"),be=a("a"),Jn=n("This blog post"),Qn=n(" provides an overview of the comparison benchmarks for XLA-compatible models along with a friendly introduction to XLA in TensorFlow."),Vn=c(),Ce=a("li"),ke=a("a"),Zn=n("This blog post"),es=n(" discusses our design philosophy behind adding XLA support to the TensorFlow models in \u{1F917} Transformers."),ts=c(),Me=a("li"),os=n("Recommended posts for learning more about XLA and TensorFlow graphs in general:"),P=a("ul"),gt=a("li"),ye=a("a"),ns=n("XLA: Optimizing Compiler for Machine Learning"),ss=c(),_t=a("li"),Te=a("a"),as=n("Introduction to graphs and tf.function"),rs=c(),wt=a("li"),Ae=a("a"),is=n("Better performance with tf.function"),this.h()},l(e){const l=ka('[data-svelte="svelte-1phssyn"]',document.head);j=r(l,"META",{name:!0,content:!0}),l.forEach(o),bt=u(e),L=r(e,"H1",{class:!0});var no=i(L);M=r(no,"A",{id:!0,class:!0,href:!0});var us=i(M);Ge=r(us,"SPAN",{});var ds=i(Ge);g(K.$$.fragment,ds),ds.forEach(o),us.forEach(o),go=u(no),Ie=r(no,"SPAN",{});var hs=i(Ie);_o=s(hs,"XLA Integration for TensorFlow Models"),hs.forEach(o),no.forEach(o),kt=u(e),g(J.$$.fragment,e),yt=u(e),O=r(e,"P",{});var so=i(O);wo=s(so,"Accelerated Linear Algebra, dubbed XLA, is a compiler for accelerating the runtime of TensorFlow Models. From the "),Q=r(so,"A",{href:!0,rel:!0});var ms=i(Q);vo=s(ms,"official documentation"),ms.forEach(o),bo=s(so,":"),so.forEach(o),Tt=u(e),xe=r(e,"P",{});var gs=i(xe);ko=s(gs,"XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes."),gs.forEach(o),At=u(e),d=r(e,"P",{});var h=i(d);yo=s(h,"Using XLA in TensorFlow is simple \u2013 it comes packaged inside the "),Ne=r(h,"CODE",{});var _s=i(Ne);To=s(_s,"tensorflow"),_s.forEach(o),Ao=s(h," library, and it can be triggered with the "),Re=r(h,"CODE",{});var ws=i(Re);$o=s(ws,"jit_compile"),ws.forEach(o),xo=s(h," argument in any graph-creating function such as "),V=r(h,"A",{href:!0,rel:!0});var vs=i(V);He=r(vs,"CODE",{});var bs=i(He);Eo=s(bs,"tf.function"),bs.forEach(o),vs.forEach(o),jo=s(h,". When using Keras methods like "),Ye=r(h,"CODE",{});var ks=i(Ye);Lo=s(ks,"fit()"),ks.forEach(o),Fo=s(h," and "),Ue=r(h,"CODE",{});var ys=i(Ue);Xo=s(ys,"predict()"),ys.forEach(o),qo=s(h,", you can enable XLA simply by passing the "),Be=r(h,"CODE",{});var Ts=i(Be);zo=s(Ts,"jit_compile"),Ts.forEach(o),Po=s(h," argument to "),Ke=r(h,"CODE",{});var As=i(Ke);Co=s(As,"model.compile()"),As.forEach(o),Mo=s(h,". However, XLA is not limited to these methods - it can also be used to accelerate any arbitrary "),Je=r(h,"CODE",{});var $s=i(Je);Oo=s($s,"tf.function"),$s.forEach(o),Do=s(h,"."),h.forEach(o),$t=u(e),k=r(e,"P",{});var E=i(k);Wo=s(E,"Several TensorFlow methods in \u{1F917} Transformers have been rewritten to be XLA-compatible, including text generation for models such as "),Z=r(E,"A",{href:!0,rel:!0});var xs=i(Z);So=s(xs,"GPT2"),xs.forEach(o),Go=s(E,", "),ee=r(E,"A",{href:!0,rel:!0});var Es=i(ee);Io=s(Es,"T5"),Es.forEach(o),No=s(E," and "),te=r(E,"A",{href:!0,rel:!0});var js=i(te);Ro=s(js,"OPT"),js.forEach(o),Ho=s(E,", as well as speech processing for models such as "),oe=r(E,"A",{href:!0,rel:!0});var Ls=i(oe);Yo=s(Ls,"Whisper"),Ls.forEach(o),Uo=s(E,"."),E.forEach(o),xt=u(e),Ee=r(e,"P",{});var Fs=i(Ee);Bo=s(Fs,"While the exact amount of speed-up is very much model-dependent, for TensorFlow text generation models inside \u{1F917} Transformers, we noticed a speed-up of ~100x. This document will explain how you can use XLA for these models to get the maximum amount of performance. We\u2019ll also provide links to additional resources if you\u2019re interested to learn more about the benchmarks and our design philosophy behind the XLA integration."),Fs.forEach(o),Et=u(e),F=r(e,"H2",{class:!0});var ao=i(F);D=r(ao,"A",{id:!0,class:!0,href:!0});var Xs=i(D);Qe=r(Xs,"SPAN",{});var qs=i(Qe);g(ne.$$.fragment,qs),qs.forEach(o),Xs.forEach(o),Ko=u(ao),Ve=r(ao,"SPAN",{});var zs=i(Ve);Jo=s(zs,"Running TF functions with XLA"),zs.forEach(o),ao.forEach(o),jt=u(e),je=r(e,"P",{});var Ps=i(je);Qo=s(Ps,"Let us consider the following model in TensorFlow:"),Ps.forEach(o),Lt=u(e),g(se.$$.fragment,e),Ft=u(e),W=r(e,"P",{});var ro=i(W);Vo=s(ro,"The above model accepts inputs having a dimension of "),Ze=r(ro,"CODE",{});var Cs=i(Ze);Zo=s(Cs,"(10, )"),Cs.forEach(o),en=s(ro,". We can use the model for running a forward pass like so:"),ro.forEach(o),Xt=u(e),g(ae.$$.fragment,e),qt=u(e),Le=r(e,"P",{});var Ms=i(Le);tn=s(Ms,"In order to run the forward pass with an XLA-compiled function, we\u2019d need to do:"),Ms.forEach(o),zt=u(e),g(re.$$.fragment,e),Pt=u(e),T=r(e,"P",{});var Oe=i(T);on=s(Oe,"The default "),et=r(Oe,"CODE",{});var Os=i(et);nn=s(Os,"call()"),Os.forEach(o),sn=s(Oe," function of the "),tt=r(Oe,"CODE",{});var Ds=i(tt);an=s(Ds,"model"),Ds.forEach(o),rn=s(Oe," is used for compiling the XLA graph. But if there\u2019s any other model function you want to compile into XLA that\u2019s also possible with:"),Oe.forEach(o),Ct=u(e),g(ie.$$.fragment,e),Mt=u(e),X=r(e,"H2",{class:!0});var io=i(X);S=r(io,"A",{id:!0,class:!0,href:!0});var Ws=i(S);ot=r(Ws,"SPAN",{});var Ss=i(ot);g(le.$$.fragment,Ss),Ss.forEach(o),Ws.forEach(o),ln=u(io),nt=r(io,"SPAN",{});var Gs=i(nt);pn=s(Gs,"Running a TF text generation model with XLA from \u{1F917} Transformers"),Gs.forEach(o),io.forEach(o),Ot=u(e),G=r(e,"P",{});var lo=i(G);fn=s(lo,"To enable XLA-accelerated generation within \u{1F917} Transformers, you need to have a recent version of "),st=r(lo,"CODE",{});var Is=i(st);cn=s(Is,"transformers"),Is.forEach(o),un=s(lo," installed. You can install it by running:"),lo.forEach(o),Dt=u(e),g(pe.$$.fragment,e),Wt=u(e),Fe=r(e,"P",{});var Ns=i(Fe);dn=s(Ns,"And then you can run the following code:"),Ns.forEach(o),St=u(e),g(fe.$$.fragment,e),Gt=u(e),I=r(e,"P",{});var po=i(I);hn=s(po,"As you can notice, enabling XLA on "),at=r(po,"CODE",{});var Rs=i(at);mn=s(Rs,"generate()"),Rs.forEach(o),gn=s(po," is just a single line of code. The rest of the code remains unchanged. However, there are a couple of gotchas in the above code snippet that are specific to XLA. You need to be aware of those to realize the speed-ups that XLA can bring in. We discuss these in the following section."),po.forEach(o),It=u(e),q=r(e,"H2",{class:!0});var fo=i(q);N=r(fo,"A",{id:!0,class:!0,href:!0});var Hs=i(N);rt=r(Hs,"SPAN",{});var Ys=i(rt);g(ce.$$.fragment,Ys),Ys.forEach(o),Hs.forEach(o),_n=u(fo),it=r(fo,"SPAN",{});var Us=i(it);wn=s(Us,"Gotchas to be aware of"),Us.forEach(o),fo.forEach(o),Nt=u(e),A=r(e,"P",{});var De=i(A);vn=s(De,"When you are executing an XLA-enabled function (like "),lt=r(De,"CODE",{});var Bs=i(lt);bn=s(Bs,"xla_generate()"),Bs.forEach(o),kn=s(De," above) for the first time, it will internally try to infer the computation graph, which is time-consuming.  This process is known as "),ue=r(De,"A",{href:!0,rel:!0});var Ks=i(ue);yn=s(Ks,"\u201Ctracing\u201D"),Ks.forEach(o),Tn=s(De,"."),De.forEach(o),Rt=u(e),R=r(e,"P",{});var co=i(R);An=s(co,"You might notice that the generation time is not fast. Successive calls of "),pt=r(co,"CODE",{});var Js=i(pt);$n=s(Js,"xla_generate()"),Js.forEach(o),xn=s(co," (or any other XLA-enabled function) won\u2019t have to infer the computation graph, given the inputs to the function follow the same shape with which the computation graph was initially built. While this is not a problem for modalities with fixed input shapes (e.g., images), you must pay attention if you are working with variable input shape modalities (e.g., text)."),co.forEach(o),Ht=u(e),$=r(e,"P",{});var We=i($);En=s(We,"To ensure "),ft=r(We,"CODE",{});var Qs=i(ft);jn=s(Qs,"xla_generate()"),Qs.forEach(o),Ln=s(We," always operates with the same input shapes, you can specify the "),ct=r(We,"CODE",{});var Vs=i(ct);Fn=s(Vs,"padding"),Vs.forEach(o),Xn=s(We," arguments when calling the tokenizer."),We.forEach(o),Yt=u(e),g(de.$$.fragment,e),Ut=u(e),H=r(e,"P",{});var uo=i(H);qn=s(uo,"This way, you can ensure that the inputs to "),ut=r(uo,"CODE",{});var Zs=i(ut);zn=s(Zs,"xla_generate()"),Zs.forEach(o),Pn=s(uo," will always receive inputs with the shape it was traced with and thus leading to speed-ups in the generation time. You can verify this with the code below:"),uo.forEach(o),Bt=u(e),g(he.$$.fragment,e),Kt=u(e),Xe=r(e,"P",{});var ea=i(Xe);Cn=s(ea,"On a Tesla T4 GPU, you can expect the outputs like so:"),ea.forEach(o),Jt=u(e),g(me.$$.fragment,e),Qt=u(e),Y=r(e,"P",{});var ho=i(Y);Mn=s(ho,"The first call to "),dt=r(ho,"CODE",{});var ta=i(dt);On=s(ta,"xla_generate()"),ta.forEach(o),Dn=s(ho," is time-consuming because of tracing, but the successive calls are orders of magnitude faster. Keep in mind that any change in the generation options at any point with trigger re-tracing and thus leading to slow-downs in the generation time."),ho.forEach(o),Vt=u(e),qe=r(e,"P",{});var oa=i(qe);Wn=s(oa,"We didn\u2019t cover all the text generation options \u{1F917} Transformers provides in this document. We encourage you to read the documentation for advanced use cases."),oa.forEach(o),Zt=u(e),z=r(e,"H2",{class:!0});var mo=i(z);U=r(mo,"A",{id:!0,class:!0,href:!0});var na=i(U);ht=r(na,"SPAN",{});var sa=i(ht);g(ge.$$.fragment,sa),sa.forEach(o),na.forEach(o),Sn=u(mo),mt=r(mo,"SPAN",{});var aa=i(mt);Gn=s(aa,"Additional Resources"),aa.forEach(o),mo.forEach(o),eo=u(e),ze=r(e,"P",{});var ra=i(ze);In=s(ra,"Here, we leave you with some additional resources if you want to delve deeper into XLA in \u{1F917} Transformers and in general."),ra.forEach(o),to=u(e),y=r(e,"UL",{});var B=i(y);x=r(B,"LI",{});var $e=i(x);_e=r($e,"A",{href:!0,rel:!0});var ia=i(_e);Nn=s(ia,"This Colab Notebook"),ia.forEach(o),Rn=s($e," provides an interactive demonstration if you want to fiddle with the XLA-compatible encoder-decoder (like "),we=r($e,"A",{href:!0,rel:!0});var la=i(we);Hn=s(la,"T5"),la.forEach(o),Yn=s($e,") and decoder-only (like "),ve=r($e,"A",{href:!0,rel:!0});var pa=i(ve);Un=s(pa,"GPT2"),pa.forEach(o),Bn=s($e,") text generation models."),$e.forEach(o),Kn=u(B),Pe=r(B,"LI",{});var ls=i(Pe);be=r(ls,"A",{href:!0,rel:!0});var fa=i(be);Jn=s(fa,"This blog post"),fa.forEach(o),Qn=s(ls," provides an overview of the comparison benchmarks for XLA-compatible models along with a friendly introduction to XLA in TensorFlow."),ls.forEach(o),Vn=u(B),Ce=r(B,"LI",{});var ps=i(Ce);ke=r(ps,"A",{href:!0,rel:!0});var ca=i(ke);Zn=s(ca,"This blog post"),ca.forEach(o),es=s(ps," discusses our design philosophy behind adding XLA support to the TensorFlow models in \u{1F917} Transformers."),ps.forEach(o),ts=u(B),Me=r(B,"LI",{});var fs=i(Me);os=s(fs,"Recommended posts for learning more about XLA and TensorFlow graphs in general:"),P=r(fs,"UL",{});var Se=i(P);gt=r(Se,"LI",{});var ua=i(gt);ye=r(ua,"A",{href:!0,rel:!0});var da=i(ye);ns=s(da,"XLA: Optimizing Compiler for Machine Learning"),da.forEach(o),ua.forEach(o),ss=u(Se),_t=r(Se,"LI",{});var ha=i(_t);Te=r(ha,"A",{href:!0,rel:!0});var ma=i(Te);as=s(ma,"Introduction to graphs and tf.function"),ma.forEach(o),ha.forEach(o),rs=u(Se),wt=r(Se,"LI",{});var ga=i(wt);Ae=r(ga,"A",{href:!0,rel:!0});var _a=i(Ae);is=s(_a,"Better performance with tf.function"),_a.forEach(o),ga.forEach(o),Se.forEach(o),fs.forEach(o),B.forEach(o),this.h()},h(){f(j,"name","hf:doc:metadata"),f(j,"content",JSON.stringify(xa)),f(M,"id","xla-integration-for-tensorflow-models"),f(M,"class","header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full"),f(M,"href","#xla-integration-for-tensorflow-models"),f(L,"class","relative group"),f(Q,"href","https://www.tensorflow.org/xla"),f(Q,"rel","nofollow"),f(V,"href","https://www.tensorflow.org/guide/intro_to_graphs"),f(V,"rel","nofollow"),f(Z,"href","https://huggingface.co/docs/transformers/model_doc/gpt2"),f(Z,"rel","nofollow"),f(ee,"href","https://huggingface.co/docs/transformers/model_doc/t5"),f(ee,"rel","nofollow"),f(te,"href","https://huggingface.co/docs/transformers/model_doc/opt"),f(te,"rel","nofollow"),f(oe,"href","https://huggingface.co/docs/transformers/model_doc/whisper"),f(oe,"rel","nofollow"),f(D,"id","running-tf-functions-with-xla"),f(D,"class","header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full"),f(D,"href","#running-tf-functions-with-xla"),f(F,"class","relative group"),f(S,"id","running-a-tf-text-generation-model-with-xla-from-transformers"),f(S,"class","header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full"),f(S,"href","#running-a-tf-text-generation-model-with-xla-from-transformers"),f(X,"class","relative group"),f(N,"id","gotchas-to-be-aware-of"),f(N,"class","header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full"),f(N,"href","#gotchas-to-be-aware-of"),f(q,"class","relative group"),f(ue,"href","https://www.tensorflow.org/guide/intro_to_graphs#when_is_a_function_tracing"),f(ue,"rel","nofollow"),f(U,"id","additional-resources"),f(U,"class","header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full"),f(U,"href","#additional-resources"),f(z,"class","relative group"),f(_e,"href","https://colab.research.google.com/github/huggingface/blog/blob/main/notebooks/91_tf_xla_generate.ipynb"),f(_e,"rel","nofollow"),f(we,"href","https://huggingface.co/docs/transformers/model_doc/t5"),f(we,"rel","nofollow"),f(ve,"href","https://huggingface.co/docs/transformers/model_doc/gpt2"),f(ve,"rel","nofollow"),f(be,"href","https://huggingface.co/blog/tf-xla-generate"),f(be,"rel","nofollow"),f(ke,"href","https://blog.tensorflow.org/2022/11/how-hugging-face-improved-text-generation-performance-with-xla.html"),f(ke,"rel","nofollow"),f(ye,"href","https://www.tensorflow.org/xla"),f(ye,"rel","nofollow"),f(Te,"href","https://www.tensorflow.org/guide/intro_to_graphs"),f(Te,"rel","nofollow"),f(Ae,"href","https://www.tensorflow.org/guide/function"),f(Ae,"rel","nofollow")},m(e,l){t(document.head,j),p(e,bt,l),p(e,L,l),t(L,M),t(M,Ge),_(K,Ge,null),t(L,go),t(L,Ie),t(Ie,_o),p(e,kt,l),_(J,e,l),p(e,yt,l),p(e,O,l),t(O,wo),t(O,Q),t(Q,vo),t(O,bo),p(e,Tt,l),p(e,xe,l),t(xe,ko),p(e,At,l),p(e,d,l),t(d,yo),t(d,Ne),t(Ne,To),t(d,Ao),t(d,Re),t(Re,$o),t(d,xo),t(d,V),t(V,He),t(He,Eo),t(d,jo),t(d,Ye),t(Ye,Lo),t(d,Fo),t(d,Ue),t(Ue,Xo),t(d,qo),t(d,Be),t(Be,zo),t(d,Po),t(d,Ke),t(Ke,Co),t(d,Mo),t(d,Je),t(Je,Oo),t(d,Do),p(e,$t,l),p(e,k,l),t(k,Wo),t(k,Z),t(Z,So),t(k,Go),t(k,ee),t(ee,Io),t(k,No),t(k,te),t(te,Ro),t(k,Ho),t(k,oe),t(oe,Yo),t(k,Uo),p(e,xt,l),p(e,Ee,l),t(Ee,Bo),p(e,Et,l),p(e,F,l),t(F,D),t(D,Qe),_(ne,Qe,null),t(F,Ko),t(F,Ve),t(Ve,Jo),p(e,jt,l),p(e,je,l),t(je,Qo),p(e,Lt,l),_(se,e,l),p(e,Ft,l),p(e,W,l),t(W,Vo),t(W,Ze),t(Ze,Zo),t(W,en),p(e,Xt,l),_(ae,e,l),p(e,qt,l),p(e,Le,l),t(Le,tn),p(e,zt,l),_(re,e,l),p(e,Pt,l),p(e,T,l),t(T,on),t(T,et),t(et,nn),t(T,sn),t(T,tt),t(tt,an),t(T,rn),p(e,Ct,l),_(ie,e,l),p(e,Mt,l),p(e,X,l),t(X,S),t(S,ot),_(le,ot,null),t(X,ln),t(X,nt),t(nt,pn),p(e,Ot,l),p(e,G,l),t(G,fn),t(G,st),t(st,cn),t(G,un),p(e,Dt,l),_(pe,e,l),p(e,Wt,l),p(e,Fe,l),t(Fe,dn),p(e,St,l),_(fe,e,l),p(e,Gt,l),p(e,I,l),t(I,hn),t(I,at),t(at,mn),t(I,gn),p(e,It,l),p(e,q,l),t(q,N),t(N,rt),_(ce,rt,null),t(q,_n),t(q,it),t(it,wn),p(e,Nt,l),p(e,A,l),t(A,vn),t(A,lt),t(lt,bn),t(A,kn),t(A,ue),t(ue,yn),t(A,Tn),p(e,Rt,l),p(e,R,l),t(R,An),t(R,pt),t(pt,$n),t(R,xn),p(e,Ht,l),p(e,$,l),t($,En),t($,ft),t(ft,jn),t($,Ln),t($,ct),t(ct,Fn),t($,Xn),p(e,Yt,l),_(de,e,l),p(e,Ut,l),p(e,H,l),t(H,qn),t(H,ut),t(ut,zn),t(H,Pn),p(e,Bt,l),_(he,e,l),p(e,Kt,l),p(e,Xe,l),t(Xe,Cn),p(e,Jt,l),_(me,e,l),p(e,Qt,l),p(e,Y,l),t(Y,Mn),t(Y,dt),t(dt,On),t(Y,Dn),p(e,Vt,l),p(e,qe,l),t(qe,Wn),p(e,Zt,l),p(e,z,l),t(z,U),t(U,ht),_(ge,ht,null),t(z,Sn),t(z,mt),t(mt,Gn),p(e,eo,l),p(e,ze,l),t(ze,In),p(e,to,l),p(e,y,l),t(y,x),t(x,_e),t(_e,Nn),t(x,Rn),t(x,we),t(we,Hn),t(x,Yn),t(x,ve),t(ve,Un),t(x,Bn),t(y,Kn),t(y,Pe),t(Pe,be),t(be,Jn),t(Pe,Qn),t(y,Vn),t(y,Ce),t(Ce,ke),t(ke,Zn),t(Ce,es),t(y,ts),t(y,Me),t(Me,os),t(Me,P),t(P,gt),t(gt,ye),t(ye,ns),t(P,ss),t(P,_t),t(_t,Te),t(Te,as),t(P,rs),t(P,wt),t(wt,Ae),t(Ae,is),oo=!0},p:ya,i(e){oo||(w(K.$$.fragment,e),w(J.$$.fragment,e),w(ne.$$.fragment,e),w(se.$$.fragment,e),w(ae.$$.fragment,e),w(re.$$.fragment,e),w(ie.$$.fragment,e),w(le.$$.fragment,e),w(pe.$$.fragment,e),w(fe.$$.fragment,e),w(ce.$$.fragment,e),w(de.$$.fragment,e),w(he.$$.fragment,e),w(me.$$.fragment,e),w(ge.$$.fragment,e),oo=!0)},o(e){v(K.$$.fragment,e),v(J.$$.fragment,e),v(ne.$$.fragment,e),v(se.$$.fragment,e),v(ae.$$.fragment,e),v(re.$$.fragment,e),v(ie.$$.fragment,e),v(le.$$.fragment,e),v(pe.$$.fragment,e),v(fe.$$.fragment,e),v(ce.$$.fragment,e),v(de.$$.fragment,e),v(he.$$.fragment,e),v(me.$$.fragment,e),v(ge.$$.fragment,e),oo=!1},d(e){o(j),e&&o(bt),e&&o(L),b(K),e&&o(kt),b(J,e),e&&o(yt),e&&o(O),e&&o(Tt),e&&o(xe),e&&o(At),e&&o(d),e&&o($t),e&&o(k),e&&o(xt),e&&o(Ee),e&&o(Et),e&&o(F),b(ne),e&&o(jt),e&&o(je),e&&o(Lt),b(se,e),e&&o(Ft),e&&o(W),e&&o(Xt),b(ae,e),e&&o(qt),e&&o(Le),e&&o(zt),b(re,e),e&&o(Pt),e&&o(T),e&&o(Ct),b(ie,e),e&&o(Mt),e&&o(X),b(le),e&&o(Ot),e&&o(G),e&&o(Dt),b(pe,e),e&&o(Wt),e&&o(Fe),e&&o(St),b(fe,e),e&&o(Gt),e&&o(I),e&&o(It),e&&o(q),b(ce),e&&o(Nt),e&&o(A),e&&o(Rt),e&&o(R),e&&o(Ht),e&&o($),e&&o(Yt),b(de,e),e&&o(Ut),e&&o(H),e&&o(Bt),b(he,e),e&&o(Kt),e&&o(Xe),e&&o(Jt),b(me,e),e&&o(Qt),e&&o(Y),e&&o(Vt),e&&o(qe),e&&o(Zt),e&&o(z),b(ge),e&&o(eo),e&&o(ze),e&&o(to),e&&o(y)}}}const xa={local:"xla-integration-for-tensorflow-models",sections:[{local:"running-tf-functions-with-xla",title:"Running TF functions with XLA"},{local:"running-a-tf-text-generation-model-with-xla-from-transformers",title:"Running a TF text generation model with XLA from \u{1F917} Transformers"},{local:"gotchas-to-be-aware-of",title:"Gotchas to be aware of"},{local:"additional-resources",title:"Additional Resources"}],title:"XLA Integration for TensorFlow Models"};function Ea(cs){return Ta(()=>{new URLSearchParams(window.location.search).get("fw")}),[]}class qa extends wa{constructor(j){super();va(this,j,Ea,$a,ba,{})}}export{qa as default,xa as metadata};
