top of page

Accelerated Linear Algebra(XLA), a JIT compiler for AI applications

Image source: Mercedes AMG F1 W10 EQ Power+ - Wikipedia

XLA is a compiler for accelerating the runtime of TensorFlow Models.


Tracing When executing an XLA-enabled function (like xla_generate() for the first time, it will internally

try to infer the computational graph, which is time-consuming. But following executions will be fast as

the graph won't be generated again

AI workload specific compiler for linear algebra that can accelerate TensorFlow models with potentially

no source code changes.

XLA comes packaged inside the tensorflow library, and it can be triggered with the jit_compile argument

in any graph-creating function such as tf.function.

When using Keras methods like fit() and predict(), you can enable XLA simply by passing the jit_compile

argument to model.compile()

Eager execution refers to a programming environment that evaluates operations immediately. Python

implements eager execution

TensorFlow 2.0 has a feature called AutoGraph which lets you write TensorFlow graph code using

regular Python code syntax by using the decorator @tf.function

Additional Reading

The TensorFlow XLA guide provides an introduction to XLA, explaining its concepts, features, and benefits. It also includes instructions on how to use XLA with TensorFlow and how to enable XLA compilation for improved performance.

TensorFlow XLA GitHub Repository

TensorFlow XLA Code Labs: TensorFlow provides code labs that offer hands-on exercises and tutorials for XLA.

bottom of page