Running TensorFlow with XLA tf.function throws error

Question

When I'm trying to compile this code, getting the below error.

File "xla_test.py", line 25, in <module>
    @tf.function(jit_compile=True)
TypeError: function() got an unexpected keyword argument 'jit_compile'

Answer 1

Without switching to tf-nightly just use:

@tf.function(experimental_compile=True)

From the tensorflow docs:

experimental_compile If True, the function is always compiled by XLA. XLA may be more efficient in some cases (e.g. TPU, XLA_GPU, dense tensor computations).

In my case MCMC sampling took without that parameter: ~1 min 37 sec, with experimental_compile=True: ~6 sec. Tensorflow built from sources (r2.4 branch).

Answer 2

Installing tf-nightly resolved this issue.

pip install tf-nightly