From f758d3e8dd4d533a46409b1cdaf55ad74deb4ec8 Mon Sep 17 00:00:00 2001 From: zpcore Date: Thu, 2 May 2024 22:34:15 +0000 Subject: [PATCH] pass TPU_ML_PLATFORM_VERSION env to libtpu --- torch_xla/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py index 6b83d45e4b4..2841a65c885 100644 --- a/torch_xla/__init__.py +++ b/torch_xla/__init__.py @@ -6,6 +6,7 @@ import torch import _XLAC from ._internal import tpu +from .version import __version__ logging.basicConfig() logger = logging.getLogger(__name__) @@ -76,6 +77,8 @@ def _setup_default_env(): os.environ.setdefault('ALLOW_MULTIPLE_LIBTPU_LOAD', '1') os.environ.setdefault('TPU_ML_PLATFORM', 'PyTorch/XLA') + # This is used for ML Framework Telemetry. + os.environ.setdefault('TPU_ML_PLATFORM_VERSION', __version__) if tpu.version() == 4: os.environ.setdefault('TPU_MEGACORE', 'megacore_dense') @@ -149,7 +152,6 @@ def _setup_tpu_vm_library_path() -> bool: import atexit from ._patched_functions import _apply_patches -from .version import __version__ _found_libtpu = _setup_tpu_vm_library_path()