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