From f24cb03585fc8ac6f81168c6f6a310755fd0a47a Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Wed, 24 Jun 2026 12:48:15 -0700 Subject: [PATCH] fix serialization of int values matching tensorflow.ProfileOptions.AdvancedConfigValue PiperOrigin-RevId: 937504146 --- pathwaysutils/profiling.py | 2 +- pathwaysutils/test/profiling_test.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 7e067bc..029fc83 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -141,7 +141,7 @@ def _create_profile_request( if isinstance(v, bool): advanced_config[k] = {"boolValue": v} elif isinstance(v, int): - advanced_config[k] = {"intValue": v} + advanced_config[k] = {"int64Value": v} elif isinstance(v, str): advanced_config[k] = {"stringValue": v} else: diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index c189fb0..d6d926b 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -569,6 +569,8 @@ def test_create_profile_request_with_options(self): "tpu_num_chips_to_profile_per_task": 3, "tpu_num_sparse_core_tiles_to_trace": 5, "tpu_trace_mode": "TRACE_COMPUTE", + "tpu_num_sparse_cores_to_trace": 1, + "tpu_enable_flag": True, } request = profiling._create_profile_request( @@ -586,9 +588,15 @@ def test_create_profile_request_with_options(self): "pwTraceOptions": { "enablePythonTracer": True, "advancedConfiguration": { - "tpu_num_chips_to_profile_per_task": {"intValue": 3}, - "tpu_num_sparse_core_tiles_to_trace": {"intValue": 5}, + "tpu_num_chips_to_profile_per_task": { + "int64Value": 3 + }, + "tpu_num_sparse_core_tiles_to_trace": { + "int64Value": 5 + }, "tpu_trace_mode": {"stringValue": "TRACE_COMPUTE"}, + "tpu_num_sparse_cores_to_trace": {"int64Value": 1}, + "tpu_enable_flag": {"boolValue": True}, }, }, },