CODEWITHSUNDEEP
tensorflow2.0tensorflow-xla
harry
harry

Reputation: 1093

Running TensorFlow with XLA tf.function throws error

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'

Upvotes: 3

Views: 2207

Answers (2)

Vlad
Vlad

Reputation: 21

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).

Upvotes: 2

harry
harry

Reputation: 1093

Installing tf-nightly resolved this issue.

pip install tf-nightly

Upvotes: 0

Related Questions