Nirmal Reddy Muppiri

@superbobry do you mean something like this should suffice? ``` def start_trace(log_dir, create_perfetto_link: bool = False, create_perfetto_trace: bool = False) -> xla_client.profiler.ProfilerSession: #....code above jax_profile_output = os.getenv('JAX_PROFILE') if jax_profile_output: log_dir...

Cool, thanks for the patience! I've updated everything with the changes, so please feel free to review/merge.