Add TensorFlow version to TPU_ML_PLATFORM_VERSION environment varia…

…bles. This will allow us to track the version of TensorFlow that is being used to run on Cloud TPU. PiperOrigin-RevId: 636219692
tensorflow · May 22, 2024 · 62d7c31 · 62d7c31
1 parent e5592ac
commit 62d7c31
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/python/tpu/__init__.py b/tensorflow/python/tpu/__init__.py
@@ -16,4 +16,7 @@
 """Ops related to Tensor Processing Units."""
 
 import os
+from tensorflow.python.framework import versions
+
 os.environ['TPU_ML_PLATFORM'] = 'Tensorflow'
+os.environ['TPU_ML_PLATFORM_VERSION'] = versions.__version__