

In order to run the forward pass with an XLA-compiled function, we’d need to do:
py
xla_fn = tf.function(model, jit_compile=True)
_ = xla_fn(random_inputs)
The default call() function of the model is used for compiling the XLA graph. But if there’s any other model function you want to compile into XLA that’s also possible with:
py
my_xla_fn = tf.function(model.my_xla_fn, jit_compile=True)
Running a TF text generation model with XLA from 🤗 Transformers
To enable XLA-accelerated generation within 🤗 Transformers, you need to have a recent version of transformers installed. You can install it by running: