From 3b8b8460c40c136c37be0887b9b7a8765c2656ed Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Mon, 8 Dec 2025 15:47:28 -0800 Subject: [PATCH] Expose profiler advanced configuration as a Python dict. In profiler.cc, the advanced_configuration property of tensorflow::ProfileOptions is now exposed as a Python dictionary. The getter converts the proto map to a nb::dict, handling different value types (bool, int64, string). Example error: ``` ProfileOptions().advanced_configuration ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ TypeError: Unable to convert function return value to a Python type! The signature was (self) -> proto2::Map, std::__u::allocator>, tensorflow::ProfileOptions_AdvancedConfigValue> ``` PiperOrigin-RevId: 841944949 --- tests/profiler_test.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/profiler_test.py b/tests/profiler_test.py index 3088ced9872a..8ab4940e1bb5 100644 --- a/tests/profiler_test.py +++ b/tests/profiler_test.py @@ -508,5 +508,18 @@ def on_profile(): unittest.mock.ANY, ) + @unittest.skipIf(jax._src.lib.ifrt_version < 39, "advanced_configuration getter is newly added") + def test_advanced_configuration_getter(self): + options = jax.profiler.ProfileOptions() + advanced_config = { + "tpu_trace_mode": "TRACE_COMPUTE", + "tpu_num_sparse_cores_to_trace": 1, + "enableFwThrottleEvent": True, + } + options.advanced_configuration = advanced_config + returned_config = options.advanced_configuration + self.assertIsInstance(returned_config, dict) + self.assertEqual(returned_config, advanced_config) + if __name__ == "__main__": absltest.main(testLoader=jtu.JaxTestLoader())