System information

• TensorFlow version (you are using): 2.5.0-dev20201115
• Are you willing to contribute it (Yes/No): No (don't know how to)

Describe the feature and the current behavior/state.
Currently, TF32 works in the normal run time but not with XLA.

Will this change the current api? How?
No.

Who will benefit with this feature?
Anyone using Ampere GPUs and XLA compilation.

Any Other info.
I checked the TF32 was not being used with XLA using the following script:

import tensorflow as tf
tf.debugging.set_log_device_placement(True)
tf.config.experimental.enable_tensor_float_32_execution(True)

@tf.function(experimental_compile=True)
def f(x):
    return x @ x
with tf.device("/GPU:0"):
    x = tf.cast(tf.random.uniform(shape=(16,16)), tf.float32)
f(x)

I ran this on an A100 with dlprof. With experimental_compile=False, the following kernel ran:

With experimental_compile=False and tf.config.experimental.enable_tensor_float_32_execution(False), the following kernel ran:

With experimental_compile=True and tf.config.experimental.enable_tensor_float_32_execution(True), the same kernel ran:

This indicates to me that the TF32 is not enabled with XLA. Please let me know if I'm doing something wrong. Thanks!