XLA (Accelerated Linear Algebra) is a domain-specific linear algebra compiler that speeds up the running of TensorFlow models, possibly with no source code changes at all. The usage is the same as that of native Tensorflow.

Enable XLA for full graph

sess_config = tf.ConfigProto()
sess_config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1


with tf.train.MonitoredTrainingSession(
    config=sess_config) as sess:

Enable XLA for subgraph

def dnn(self, dnn_input, dnn_hidden_units=None, layer_name=''):
    # Add the following code to limit part of the graph to use XLA to compile.
    jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
    with jit_scope():
       for layer_id, num_hidden_units in enumerate(dnn_hidden_units):
            with tf.variable_scope(layer_name + "_%d" % layer_id,
                                   reuse=tf.AUTO_REUSE) as dnn_layer_scope:
                dnn_input = tf.layers.dense(dnn_input,
                if self.use_bn:
                    dnn_input = tf.layers.batch_normalization(
                        dnn_input, training=self._is_training, trainable=True)
                add_layer_summary(dnn_input, dnn_layer_scope.name)

       return dnn_input