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

Don't compute the log message for checkpointing residuals unless it is going to be logged.

Browse Source 
PiperOrigin-RevId: 606811594
This commit is contained in:
Peter Hawkins 2024-02-13 18:23:02 -08:00 committed by jax authors
parent ac45b8f73d
commit 7156f20b44
1 changed files with 18 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/_src/ad_checkpoint.py
View File

@ -552,13 +552,14 @@ def remat_partial_eval(trace, *tracers, jaxpr, **params):
source_info_util.current())
# log info about saved residuals
log_level = logging.WARNING if config.log_checkpoint_residuals.value else logging.DEBUG
if logger.isEnabledFor(log_level):
try:
_, staged_unk = partition_list(in_used_staged, in_unknowns)
res_invars, _ = partition_list(staged_unk, jaxpr_unknown.invars[num_res:])
res_outvars = jaxpr_known.outvars[len(jaxpr_known.outvars) - num_res:]
body_res = _saved_residuals(jaxpr_known.replace(outvars=res_outvars), None)
logger.log(logging.WARNING if config.log_checkpoint_residuals.value
else logging.DEBUG,
logger.log(log_level,
'remat-decorated function ' +
'saving inputs with shapes:\n' * bool(res_invars) +
' %s\n' * len(res_invars) +