Merge pull request #21289 from jakevdp:stacklevel

PiperOrigin-RevId: 634895282
google · May 17, 2024 · 7a3fc71 · 7a3fc71
2 parents 56301d4 + 9ad6729
commit 7a3fc71
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
@@ -1224,7 +1224,7 @@ def wrap_with_deprecation_warning(f):
          "`from jax.experimental.export import export` you should use "
          "`from jax.experimental import export`.")
   def wrapped_f(*args, **kwargs):
-    warnings.warn(msg, DeprecationWarning)
+    warnings.warn(msg, DeprecationWarning, stacklevel=2)
     return f(*args, **kwargs)
   return wrapped_f