Skip to content

Commit f2f287e

Browse files
lukebaumannGoogle-ML-Automation
authored andcommitted
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::basic_string<char, std::__u::char_traits<char>, std::__u::allocator<char>>, tensorflow::ProfileOptions_AdvancedConfigValue> ``` PiperOrigin-RevId: 841944949
1 parent 0339ad1 commit f2f287e

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

tests/profiler_test.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,5 +508,17 @@ def on_profile():
508508
unittest.mock.ANY,
509509
)
510510

511+
def test_advanced_configuration_getter(self):
512+
options = jax.profiler.ProfileOptions()
513+
advanced_config = {
514+
"tpu_trace_mode": "TRACE_COMPUTE",
515+
"tpu_num_sparse_cores_to_trace": 1,
516+
"enableFwThrottleEvent": True,
517+
}
518+
options.advanced_configuration = advanced_config
519+
returned_config = options.advanced_configuration
520+
self.assertIsInstance(returned_config, dict)
521+
self.assertEqual(returned_config, advanced_config)
522+
511523
if __name__ == "__main__":
512524
absltest.main(testLoader=jtu.JaxTestLoader())

0 commit comments

Comments
 (0)