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

tfp.experimental.nn.util.tfcompile

View source on GitHub

Centralizes TF compilation related options.

tfp.experimental.nn.util.tfcompile(
    func=None, tf_function=True, xla_best_effort=True, xla_compile_all=False
)

Args:

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

Returns:

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

Example Usage

tfn = tfp.experimental.nn

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

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