[XLA:GPU] Add option to return FDO profile as textproto.

PiperOrigin-RevId: 619105468
sshahrokhi · Mar 26, 2024 · 33cf53c · 33cf53c
1 parent 86a7086
commit 33cf53c
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/jax/_src/profiler.py b/jax/_src/profiler.py
@@ -24,8 +24,7 @@
 import os
 import socketserver
 import threading
-
-from typing import Callable
+from typing import Callable, Union
 
 from jax._src import traceback_util
 traceback_util.register_exclusion(__file__)
@@ -211,7 +210,7 @@ def stop_trace():
     _profile_state.reset()
 
 
-def stop_and_get_fdo_profile() -> bytes:
+def stop_and_get_fdo_profile() -> Union[bytes, str]:
   """Stops the currently-running profiler trace and export fdo_profile.
 
   Currently, this is only supported for GPU.