Missed TensorFlow Dev Summit? Check out the video playlist. Watch recordings


View source on GitHub

Centralizes TF compilation related options.

    func=None, tf_function=True, xla_best_effort=True, xla_compile_all=False


  • func: Python callable to wrapped with the specified TF compilation directives. Default value: None.
  • tf_function: bool representing whether the resulting function should be tf.function decoreated. Default value: True.
  • xla_best_effort: bool representing whether XLA auto-clustering compilation should be performed. (This argument is ignored if the function is executed eagerly.) Default value: True.
  • xla_compile_all: bool representing whether XLA compilation should be performed. (This argument overrides both tf_function and xla_best_effort. Default value: False.


  • wrapped_func: A Python callable with the specified compilation directives embedded.

Example Usage

tfn = tfp.experimental.nn

# Use style #1.
def foo(...):

# Use style #2.
def foo(...):
foo = tfn.util.tfcompile(xla_compile_all=True)(foo)