Fix return type and tolerance in test

runame · runame · commit a3cbcd114a14 · 2025-02-23T23:57:21.000Z
diff --git a/distributed_shampoo/utils/shampoo_preconditioner_list.py b/distributed_shampoo/utils/shampoo_preconditioner_list.py
@@ -1165,7 +1165,7 @@ def _adaptive_amortized_computation_frequency_criterion_below_or_equal_tolerance
         off_diagonal_summed = squared_approximate_eigenvalues.fill_diagonal_(0.0).sum()
         norm = torch.sqrt(diagonal_summed + off_diagonal_summed)
         off_diagonal_norm = torch.sqrt(off_diagonal_summed)
-        return off_diagonal_norm <= tolerance * norm
+        return bool(off_diagonal_norm <= tolerance * norm)
 
     @torch.compiler.disable
     def _amortized_computation(self) -> None:
diff --git a/distributed_shampoo/utils/tests/shampoo_preconditioner_list_test.py b/distributed_shampoo/utils/tests/shampoo_preconditioner_list_test.py
@@ -1136,7 +1136,7 @@ def test_adaptive_amortized_computation_frequency_criterion_below_tolerance(
         assert not (off_diagonal_norm <= test_tolerance1 * norm)
 
         # Below tolerance.
-        test_tolerance2 = 2e-7
+        test_tolerance2 = 3e-7
         test_criterion2 = EigenvalueCorrectedShampooPreconditionerList._adaptive_amortized_computation_frequency_criterion_below_or_equal_tolerance(
             test_factor_matrix, test_factor_matrix_eigenvectors, test_tolerance2
         )

Original file line number	Diff line number	Diff line change
`@@ -1136,7 +1136,7 @@ def test_adaptive_amortized_computation_frequency_criterion_below_tolerance(`
`1136`	`1136`	`assert not (off_diagonal_norm <= test_tolerance1 * norm)`
`1137`	`1137`
`1138`	`1138`	`# Below tolerance.`
`1139`		`- test_tolerance2 = 2e-7`
	`1139`	`+ test_tolerance2 = 3e-7`
`1140`	`1140`	`test_criterion2 = EigenvalueCorrectedShampooPreconditionerList._adaptive_amortized_computation_frequency_criterion_below_or_equal_tolerance(`
`1141`	`1141`	`test_factor_matrix, test_factor_matrix_eigenvectors, test_tolerance2`
`1142`	`1142`	`)`