backward amp_scale

pomonam · Jul 11, 2024 · 4adc1b7 · 4adc1b7
1 parent 7220ecc
commit 4adc1b7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/kronfluence/arguments.py b/kronfluence/arguments.py
@@ -151,6 +151,10 @@ def __post_init__(self) -> None:
         ):
             raise ValueError("All data and module partitions must be positive.")
 
+        # For backward compatibility:
+        if not hasattr(self, "amp_scale"):
+            self.amp_scale = 2.0**16
+
 
 @dataclass
 class ScoreArguments(Arguments):

diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py
@@ -104,6 +104,7 @@ def test_default_factor_arguments() -> None:
     assert factor_args.strategy == "ekfac"
     assert factor_args.use_empirical_fisher is False
     assert factor_args.amp_dtype is None
+    assert factor_args.amp_scale == 2.0**16
     assert factor_args.has_shared_parameters is False
 
     assert factor_args.covariance_max_examples == 100_000