Visualize graphs before and after Grappler passes?

[Paul Delestrac]

Hi, 

I've been trying to visualize the graph of a `tf.function` with and without Grappler optimizations but so far I’m not managing to see any difference in the generated graphs. 

Here is the process I followed: I took the code used in the “TensorFlow graph optimization with Grappler” tutorial to visualize the graph with and without constant folding. Then, I use the `experimental_get_compiler_ir` function with the “hlo” setting to generate the graph before XLA optimizations: `simple_function.experimental_get_compiler_ir(x)(stage="hlo")`. Finally, I use the `interactive_graphviz` XLA tool to visualize the graphs. Unfortunately, I’m not able to see any difference between the generated graphs: in both cases constant folding has not been optimized.

When designing the mentioned process, I was assuming that Grappler optimizations are executed first and then XLA optimizations take over. However, I just found a recent post suggesting that the use of XLA could disable Grappler passes, which doesn’t make sense to me but it’s consistent with my observation.

Is there anything wrong in the process I’m following? If so, is there an alternative way to visualize graphs before and after Grappler passes?

(code used to generate the graphs in attachment)