Module: tfp.experimental.auto_batching.xla

View source on GitHub

XLA utilities.

Functions

compile_nested_output(...): Wraps f with a tpu.rewrite or xla.compile, propagates output structure.