ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tfp.experimental.nn.util.tfcompile

Centralizes TF compilation related options.

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)