mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[AutoPGLE] Update fdo_profile comment.

Browse Source 
PiperOrigin-RevId: 700759386
This commit is contained in:
jax authors 2024-11-27 11:23:34 -08:00
parent df6758f021
commit c2c177eee8
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jax/_src/pjit.py
View File

@ -1662,10 +1662,8 @@ def _pjit_call_impl_python(
pgle_compile_options['fdo_profile'] = fdo_profile
compiler_options_kvs = compiler_options_kvs + tuple(pgle_compile_options.items())
# TODO(patrios): Do not pass mutable profile session through cached lowering
# chain. Instead we need to move profilers dictionary to pxla module and use
# module as key. Right now we can't do that since there is no way to evict
# _pjit_lower_cached cache for in PGLE mode.
# Passing mutable PGLE profile here since it should be extracted by JAXPR to
# initialize the fdo_profile compile option.
compiled = _resolve_and_lower(
args, jaxpr=jaxpr, in_shardings=in_shardings,
out_shardings=out_shardings, in_layouts=in_layouts,