Filter some decorator from traceback

PiperOrigin-RevId: 733039366
google-research · Mar 3, 2025 · eb5cc17 · eb5cc17
1 parent 939b2f6
commit eb5cc17
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/kauldron/typing/type_check.py b/kauldron/typing/type_check.py
@@ -122,6 +122,9 @@ def typechecked(fn):
   @jaxtyping.jaxtyped(typechecker=None)
   @functools.wraps(fn)
   def _reraise_with_shape_info(*args, _typecheck: bool = True, **kwargs):
+    # Hide the function from the traceback. Supported by Pytest and IPython
+    __tracebackhide__ = True  # pylint: disable=unused-variable,invalid-name
+
     if not (TYPECHECKING_ENABLED and _typecheck):
       # typchecking disabled globally or locally -> just return fn(...)
       return fn(*args, **kwargs)

diff --git a/kauldron/utils/train_property.py b/kauldron/utils/train_property.py
@@ -130,6 +130,9 @@ def _add_is_training_kwargs(fn: _FnT) -> _FnT:
 
   @_internal.wraps_with_reload(fn)
   def decorated(*args, is_training_property: bool | None = None, **kwargs):  # pylint: disable=redefined-outer-name
+    # Hide the function from the traceback. Supported by Pytest and IPython
+    __tracebackhide__ = True  # pylint: disable=unused-variable,invalid-name
+
     if is_training_property is not None:
       cm = _set_train_property(is_training_property)
     else: