grappler with and without XLA

[Moshe Maor]

Hi,

I noticed that when XLA is used for inference for a tf.function, grappler is not being executed. is this the expected behavior? why is that? 

example script: 

import tensorflow as tf

a = tf.reshape(tf.range(1, 28, dtype=float),[3,3,3])

b = tf.reshape(tf.range(10, -17, -1, dtype=float),[3,3,3])

--> either @tf.function(jit_compile=True)  or @tf.function

def myfunc(a, b):

    c = a + b

    d = c * b    

    e = d / a

    return e

print(myfunc(a, b))

this script executes with grappler when only using @tf.function but when adding jit_compile=True, grappler is not running. 

Thanks, 

Moshe 

[Jacques Pienaar]

Hey,

Yes expected. There are no graph optimization passes run with jit_compile=True at TF level, only XLA optimizations used post conversion/during compilation.

-- Jacques

--
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/e3d0eed6-9905-4d6a-9fe8-de6362a27697n%40googlegroups.com.

[t kevin]

With tf.function, tf runs in graph mode and the graph is optimized by grappler.
With jit compile = true, tf graph is converted to xla hlo first then
optimized with hlo optimization passes.

Best regards
Kevin

Moshe Maor <moshe...@gmail.com> 于2021年6月8日周二 下午11:59写道：

[Moshe Maor]

a couple of follow up questions: 
1. on a keras model.fit with XLA enabled, we get both grappler and XLA, with the XLA passes integrated into grappler. why is the difference between model.fit flow and manual jit_compile ? 

2. what is the execution path of a jit_compile that translates the graph into XLA/hlo? 

thanks, 

Moshe 

[George Karpenkov]

Hi Moshe,

This is actually an implementation wart, you will only get grappler passes if your compiled tf.function is wrapped by another uncompiled tf.function.

Normally, this should not matter, as all performance optimizations should be done by XLA, and Grappler ones are usually only interfering,

but it does cause an annoying inconsistency.

George

To view this discussion on the web visit https://groups.google.com/d/msgid/xla-dev/881d1db6-ce25-4894-b865-8a85967e071bn%40googlegroups.com.