Train Model with XLA JIT Enabled

Run the Next Cell to Show XLA JIT Training Code



In [ ]:

    
cat /root/src/main/python/xla/train_linear_xla.py

Run the Next Cell

Note: The cell below will report Aborted (core dumped), but it still generate some useful info.

Click on the following generated files in File Browser in Left Navigation

train_linear_xla.log (Log File)
hlo_graph_* (JIT Visualizations)

Reference of All XLA Operations

https://www.tensorflow.org/performance/xla/operation_semantics



In [ ]:

    
%%bash

python /root/src/main/python/xla/train_linear_xla.py &> train_linear_xla.log

dot -T png /tmp/hlo_graph_1.*.dot -o /root/notebooks/hlo_graph_1.png &>/dev/null 

dot -T png /tmp/hlo_graph_10.*.dot -o /root/notebooks/hlo_graph_10.png &>/dev/null 

dot -T png /tmp/hlo_graph_25.*.dot -o /root/notebooks/hlo_graph_25.png &>/dev/null 

dot -T png /tmp/hlo_graph_50.*.dot -o /root/notebooks/hlo_graph_50.png &>/dev/null



In [ ]: