Commit 0472116d authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

[TF:XLA] Make tf_cnn_benchmarks run on CPU with XLA.

Adds _cpu_jit to tf_cnn_benchmarks_xla BUILD rule and fixes an issue in XLA bridge triggered by XLA CPU compilation of whole graphs. In particular, modifies mark_for_compilation_pass.cc to skip _Retval nodes when looking for compilation candidates in the top level function. _Retval nodes are introduced in the input subgraph as a replacement for fetches. Including _Retval nodes into XLA clusters confuses encapsulate subgraph pass that expects a graph with no pre-existing _Retval nodes.

PiperOrigin-RevId: 177518178
parent f69915dc
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment