XLA: Cannot get IR on TPU

[hjmus]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab environment
• TensorFlow installed from (source or binary): Google Colab environment
• TensorFlow version (use command below): 2.5.0
• Python version: 3.7

Describe the current behavior
Code is in https://github.com/hjmus/mock/blob/master/TPUs_in_Colab.ipynb. I was trying to get HLO IR when running tf.function on TPU, by using experimental_get_compiler_ir(), but run into the following error:

ValueError: No matching device found for '/job:worker/replica:0/task:0/device:TPU:1'

Describe the expected behavior
Be able to get IR when running on TPU.

Standalone code to reproduce the issue
https://github.com/hjmus/mock/blob/master/TPUs_in_Colab.ipynb

Other info / logs
The purpose of this exercise is to check out the IR generated for GSPMD (https://arxiv.org/pdf/2105.04663.pdf).

#comp:xla #XLA

[tilakrayal]

I was able to reproduce the issue in tf v 2.5 and nightly,in v2.4 i am facing different error.Please find the gist of it here.

[Saduf2019]

@hjmus
Can you please refer to this comment and let us know if it helps.

[hjmus]

@Saduf2019

Hi Saduf2019, I was already using the suggested API:

try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

tf.config.experimental_connect_to_cluster(tpu)  # <<= here
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)