fix: strip _TRAINCHECK_ prefix when displaying internal tensor-tracking attrs

Add _display_attr_name() helper that maps '_TRAINCHECK_grad_ID' -> 'grad' etc.
Use it in APIContainRelation.to_display_name (removing the return-None guard)
and ConsistencyRelation.to_display_name. Remove the now-unnecessary
[internal tracking] fallback from _format_invariant_label.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Essoz

traincheck/invariant/base_cls.py

@@ -1944,6 +1944,17 @@ def _short_api_name(full_name: str) -> str:
     return ".".join(parts[-2:]) if len(parts) >= 2 else full_name
 
 
+def _display_attr_name(attr_name: str) -> str:
+    """Strip TrainCheck-internal proxy bookkeeping prefix/suffix for display.
+
+    '_TRAINCHECK_grad_ID' → 'grad'
+    'dtype' → 'dtype'  (unchanged)
+    """
+    if attr_name.startswith("_TRAINCHECK_") and attr_name.endswith("_ID"):
+        return attr_name[len("_TRAINCHECK_") : -len("_ID")]
+    return attr_name
+
+
 def read_inv_file(file_path: str | list[str]) -> list[Invariant]:
     if isinstance(file_path, str):
         file_path = [file_path]

traincheck/invariant/consistency_relation.py

@@ -14,6 +14,7 @@
     Param,
     Relation,
     VarTypeParam,
+    _display_attr_name,
 )
 from traincheck.invariant.precondition import find_precondition
 from traincheck.onlinechecker.utils import Checker_data, set_meta_vars_online
@@ -118,7 +119,7 @@ def to_display_name(params: list[Param]) -> str | None:
         if not isinstance(p, VarTypeParam):
             return None
         var_short = p.var_type.split(".")[-1]
-        attr = p.attr_name
+        attr = _display_attr_name(p.attr_name)
         return f"{var_short}.{attr} stays consistent across training steps"
 
     @staticmethod

traincheck/invariant/contain_relation.py

@@ -23,6 +23,7 @@
     Relation,
     VarNameParam,
     VarTypeParam,
+    _display_attr_name,
     _short_api_name,
     calc_likelihood,
     construct_api_param,
@@ -358,10 +359,7 @@ def to_display_name(params: list[Param]) -> str | None:
             child_short = _short_api_name(child.api_full_name)
             return f"{parent_short}() always calls {child_short}()"
         if isinstance(child, (VarTypeParam, VarNameParam)):
-            attr = child.attr_name
-            # Skip internal TrainCheck proxy bookkeeping attributes
-            if attr.startswith("_TRAINCHECK_"):
-                return None
+            attr = _display_attr_name(child.attr_name)
             var_short = child.var_type.split(".")[-1]
             pre = child.pre_value
             post = child.post_value

traincheck/reporting/checker_report.py

@@ -7,28 +7,12 @@
 from typing import Iterable
 
 from traincheck.invariant import CheckerResult, Invariant
-from traincheck.invariant.base_cls import APIParam, VarNameParam, VarTypeParam
 
 
 def _format_invariant_label(invariant: Invariant) -> str:
     display = invariant.relation.to_display_name(invariant.params)
     if display:
         return display
-    # When to_display_name returns None, fall back — but sanitize params that
-    # contain internal TrainCheck proxy bookkeeping names (_TRAINCHECK_*) so
-    # they never surface raw in the UI.
-    for p in invariant.params:
-        if isinstance(p, (VarTypeParam, VarNameParam)) and p.attr_name.startswith(
-            "_TRAINCHECK_"
-        ):
-            # Build a minimal label using only the API param, hiding internals
-            api_parts = [q for q in invariant.params if isinstance(q, APIParam)]
-            from traincheck.invariant.base_cls import _short_api_name
-
-            if api_parts:
-                func = _short_api_name(api_parts[0].api_full_name)
-                return f"{func}() [internal tracking]"
-            return f"{invariant.relation.__name__} [internal tracking]"
     if invariant.text_description:
         return invariant.text_description
     params = ", ".join(str(param) for param in invariant.params)
@@ -290,7 +274,9 @@ def build_online_report_data(
         and warm_up_steps is not None
         and sampling_interval > 0
     ):
-        checked_steps = max(0, current_step - warm_up_steps) // sampling_interval + 1
+        checked_steps = max(0, current_step - warm_up_steps) // sampling_interval + min(
+            current_step, warm_up_steps
+        )
 
     def _make_entry(inv: Invariant, count: int) -> dict:
         detail = violation_details.get(inv, {})