tfp.experimental.nn.util.tfcompile

Centralizes TF compilation related options.

tfp.experimental.nn.util.tfcompile(
    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.
@tfn.util.tfcompile(xla_compile_all=True)
def foo(...):
     ...

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