Golden test tolerance: strict by default, HIGHJAX_TESTS_LOOSE=1 for CI

Author: cool-RR

.github/workflows/tests.yml

@@ -26,6 +26,8 @@ jobs:
           pip install -e ".[trainer,tests]"
 
       - name: Run Python tests
+        env:
+          HIGHJAX_TESTS_LOOSE: '1'
         run: |
           pytest -n auto
 

tests/test_highjax_trainer/test_golden_runs/common.py

@@ -1,8 +1,14 @@
 '''Shared logic for golden run tests.'''
 from __future__ import annotations
 
+import os
+
 import jax.numpy as jnp
 
+_LOOSE = os.environ.get('HIGHJAX_TESTS_LOOSE') == '1'
+_RTOL = 1e-02 if _LOOSE else 1e-04
+_ATOL = 1e-05 if _LOOSE else 1e-07
+
 
 def print_golden_data(epoch_metrics: list[dict]) -> None:
     n = len(epoch_metrics)
@@ -38,7 +44,7 @@ def check_golden_run(epoch_metrics: list[dict], golden_data: dict) -> None:
             expected = expected_values[i]
             actual = float(metrics[key])
 
-            if not jnp.isclose(actual, expected, rtol=1e-02, atol=1e-05):
+            if not jnp.isclose(actual, expected, rtol=_RTOL, atol=_ATOL):
                 bad_fields.append((key, actual, expected))
 
         if bad_fields: