Searched refs:outfeed_rewriter (Results 1 – 2 of 2) sorted by relevance
194 outfeed_rewriter: Optional[Callable[[core.Jaxpr], core.Jaxpr]] = None variable196 if outfeed_rewriter is not None:197 return outfeed_rewriter(jaxpr)
1414 xla.outfeed_rewriter = lambda j: _rewrite_jaxpr(j, False, False)