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

Module: tfp.experimental.auto_batching.xla

XLA utilities.

Functions

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