Breaking of XLA Computation

[Naveen Swamy]

I am trying to understand how TF/XLA groups the operations into multiple XLA Computations,  basically our backend expects a single HLO of all operations. I am wondering if there is a way to do this through a flag. I trying backtracking through the code, tf python code is a tough maze, I couldn't figure out when it decides to group the operations. 

Appreciate your input.

Thanks, Naveen

[George Karpenkov]

hi Naveen, 

How do you use XLA from Python?

George

--
You received this message because you are subscribed to the Google Groups "XLA development" group.
To unsubscribe from this group and stop receiving emails from it, send an email to xla-dev+u...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/xla-dev/e7b0b700-dcb5-4b27-b16a-66c5b788d199n%40googlegroups.com.

[George Karpenkov]

Using jit_compile=True will create a single HloModule for the compiled block. A single HloModule does have multiple HloComputations though.

[Naveen Swamy]

Hi George,

Thanks for your response. We are using session to use XLA device. We are on Tensorflow 2.2

Thanks, Naveen

[Naveen Swamy]

and tf.device 

[George Karpenkov]

Could you provide a code snippet of how do you use it?