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

tree_util: fix warning category and stacklevel

Browse Source
This commit is contained in:
Jake VanderPlas 2022-07-15 09:24:22 -07:00
parent 10720258ea
commit 6907dfad00
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/tree_util.py
View File

@ -527,7 +527,8 @@ def _deprecate(f):
@functools.wraps(f)
def wrapped(*args, **kwargs):
warnings.warn(f"jax.{f.__name__} is deprecated, and will be removed in a future release. "
f"Use jax.tree_util.{f.__name__} instead.")
f"Use jax.tree_util.{f.__name__} instead.",
category=FutureWarning, stacklevel=2)
return f(*args, **kwargs)
return wrapped