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

Checkify: Remove stray raise_as_much_as_possible.

Browse Source 
We switched to initial style, and no longer need to rely on raising all tracers
to an active Checkify trace.

PiperOrigin-RevId: 510450683
This commit is contained in:
Lena Martens 2023-02-17 09:28:27 -08:00 committed by jax authors
parent edff87eb07
commit 307288b943
1 changed files with 0 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/checkify.py
View File

@ -1127,7 +1127,6 @@ def _check(pred, msg, debug, *fmt_args, **fmt_kwargs):
_check_error(error, debug=debug)
def _check_error(error, *, debug=False):
error = tree_map(core.raise_as_much_as_possible, error)
if any(map(np.shape, error._pred.values())):
error = _reduce_any_error(error)
err_args, tree_def = tree_flatten(error)