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

Merge pull request #11812 from mattjj:fix-ad-checkpoint-traceback-exclusion

Browse Source 
PiperOrigin-RevId: 466444604
This commit is contained in:
jax authors 2022-08-09 13:07:56 -07:00
commit 22a4928a74
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/ad_checkpoint.py
View File

@ -27,12 +27,14 @@ from jax.interpreters import mlir
from jax.tree_util import tree_flatten, tree_unflatten
from jax._src import ad_util
from jax._src import source_info_util
from jax._src import traceback_util
from jax._src.api_util import flatten_fun, shaped_abstractify
from jax._src.traceback_util import api_boundary
from jax._src.util import (unzip2, wraps, split_list, partition_list, safe_map,
safe_zip, merge_lists, weakref_lru_cache)
source_info_util.register_exclusion(__file__)
traceback_util.register_exclusion(__file__)
# TODO(mattjj): before this can be the standard remat implementation, we must:
# [ ] fix up callers who use the 'concrete' option (now removed)