TensorFlow: XLA not running with "Invalid argument: No _XlaCompile" errors

Question

I'm trying out the XLA tutorial described here: https://www.tensorflow.org/performance/xla/jit

I'm running mnist_softmax_xla.py from https://raw.githubusercontent.com/tensorflow/tensorflow/r1.1/tensorflow/examples/tutorials/mnist/mnist_softmax_xla.py with the following options:

TF_CPP_MIN_VLOG_LEVEL=2 TF_XLA_FLAGS='--xla_generate_hlo_graph=.*' python mnist_softmax_xla.py

Unfortunately I get a bunch of "Custom creator error: Invalid argument: No _XlaCompile for Const" errors in the output (and same for every other type of op). Also, no hlo_graph_xx.dot files are created (as the tutorial says should be).

My python install is Anaconda 4.3.1 (Anaconda3-4.3.1-Linux-x86_64.sh) on Ubuntu 16.04 LTS.

TensorFlow is version 1.1.0 compiled from source with the following commands:

$ echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
$ curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
$ sudo apt-get update && sudo apt-get install bazel

$ git clone https://github.com/tensorflow/tensorflow
$ cd tensorflow/
$ git checkout v1.1.0
$ ./configure 

Please specify the location of python. [Default is /home/ubuntu/anaconda3/bin/python]: 
Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native]: 
Do you wish to use jemalloc as the malloc implementation? [Y/n] 
jemalloc enabled
Do you wish to build TensorFlow with Google Cloud Platform support? [y/N] 
No Google Cloud Platform support will be enabled for TensorFlow
Do you wish to build TensorFlow with Hadoop File System support? [y/N] 
No Hadoop File System support will be enabled for TensorFlow
Do you wish to build TensorFlow with the XLA just-in-time compiler (experimental)? [y/N] y
XLA JIT support will be enabled for TensorFlow
Found possible Python library paths:
  /home/ubuntu/anaconda3/lib/python3.6/site-packages
Please input the desired Python library path to use.  Default is [/home/ubuntu/anaconda3/lib/python3.6/site-packages]

Using python library path: /home/ubuntu/anaconda3/lib/python3.6/site-packages
Do you wish to build TensorFlow with OpenCL support? [y/N] 
No OpenCL support will be enabled for TensorFlow
Do you wish to build TensorFlow with CUDA support? [y/N] 
No CUDA support will be enabled for TensorFlow
Configuration finished
............
INFO: Starting clean (this may take a while). Consider using --expunge_async if the clean takes more than several minutes.
...........
INFO: All external dependencies fetched successfully.

$ bazel build --config=opt //tensorflow/tools/pip_package:build_pip_package
$ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
$ pip install /tmp/tensorflow_pkg/tensorflow-1.1.0-cp36-cp36m-linux_x86_64.whl 

Why does XLA not work with this setup?

How do I install TensorFlow with working XLA?

Answer 1

Solved!

The reason for this is a little obscure note in the XLA tutorial:

Note: Turning on JIT at the session level will not result in operations being compiled for the CPU. JIT compilation for CPU operations must be done via the manual method documented below. This decision was made due to the CPU backend being single-threaded.

This doesn't work with CPU:

config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(config=config)
...

This works on CPU (and GPU as well):

jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
with jit_scope():
    ...