Add allclose for compile tests

pomonam · Jun 24, 2024 · 1a7ecfc · 1a7ecfc
1 parent fa67826
commit 1a7ecfc
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/tests/factors/test_eigens.py b/tests/factors/test_eigens.py
@@ -278,13 +278,14 @@ def test_lambda_matrices_iterative_aggregate(
     train_size: int,
     seed: int,
 ) -> None:
-    # Makes sure aggregated lambda computation is working properly.
+    # Makes sure iterative lambda computation is working properly.
     model, train_dataset, _, data_collator, task = prepare_test(
         test_name=test_name,
         train_size=train_size,
         seed=seed,
     )
     kwargs = DataLoaderKwargs(collate_fn=data_collator)
+    model = model.to(dtype=torch.float64)
     model, analyzer = prepare_model_and_analyzer(
         model=model,
         task=task,
@@ -305,13 +306,6 @@ def test_lambda_matrices_iterative_aggregate(
         factors_name=factors_name,
     )
 
-    factor_args = FactorArguments(
-        use_empirical_fisher=True,
-        lambda_iterative_aggregate=True,
-        activation_covariance_dtype=torch.float64,
-        gradient_covariance_dtype=torch.float64,
-        lambda_dtype=torch.float64,
-    )
     factor_args.lambda_iterative_aggregate = True
     analyzer.fit_all_factors(
         factors_name=factors_name + "_iterative",

diff --git a/tests/gpu_tests/compile_test.py b/tests/gpu_tests/compile_test.py
@@ -11,7 +11,8 @@
 from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
-    LAMBDA_FACTOR_NAMES,
+    LAMBDA_FACTOR_NAMES, ACTIVATION_COVARIANCE_MATRIX_NAME, GRADIENT_COVARIANCE_MATRIX_NAME,
+    NUM_ACTIVATION_COVARIANCE_PROCESSED, NUM_GRADIENT_COVARIANCE_PROCESSED,
 )
 from tests.gpu_tests.pipeline import GpuTestTask, construct_test_mlp, get_mnist_dataset
 from tests.gpu_tests.prepare_tests import QUERY_INDICES, TRAIN_INDICES
@@ -64,11 +65,20 @@ def test_covariance_matrices(self) -> None:
         )
         new_covariance_factors = self.analyzer.load_covariance_matrices(factors_name=NEW_FACTOR_NAME)
 
-        for name in COVARIANCE_FACTOR_NAMES:
+        new_covariance_factors[ACTIVATION_COVARIANCE_MATRIX_NAME] /= new_covariance_factors[NUM_ACTIVATION_COVARIANCE_PROCESSED]
+        new_covariance_factors[GRADIENT_COVARIANCE_MATRIX_NAME] /= new_covariance_factors[NUM_GRADIENT_COVARIANCE_PROCESSED]
+
+        for name in [ACTIVATION_COVARIANCE_MATRIX_NAME, GRADIENT_COVARIANCE_MATRIX_NAME]:
             for module_name in covariance_factors[name]:
                 print(f"Name: {name, module_name}")
                 print(f"Previous factor: {covariance_factors[name][module_name]}")
                 print(f"New factor: {new_covariance_factors[name][module_name]}")
+            assert check_tensor_dict_equivalence(
+                covariance_factors[name],
+                new_covariance_factors[name],
+                atol=1e-5,
+                rtol=1e-3,
+            )
 
     def test_lambda_matrices(self):
         lambda_factors = self.analyzer.load_lambda_matrices(factors_name=OLD_FACTOR_NAME)