[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

tests/ekfac_tests/compute_ekfac_ground_truth.py

@@ -20,7 +20,6 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
-import torch.distributed as dist
 import torch.nn.functional as F
 from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset
 from ground_truth.collector import (

tests/ekfac_tests/conftest.py

@@ -261,17 +261,18 @@ def ekfac_results_path(
     test_dir: str,
     ground_truth_path: str,
     ground_truth_setup: dict[str, Any],
-    overwrite: bool
+    overwrite: bool,
 ) -> str:
     """Run EKFAC computation and return results path.
 
     Uses the same data and batches as ground truth via collect_hessians to ensure
     identical batch composition and floating-point accumulation order.
     """
     import torch
+
     from bergson.config import HessianConfig
-    from bergson.hessians.hessian_approximations import collect_hessians
     from bergson.hessians.eigenvectors import compute_eigendecomposition
+    from bergson.hessians.hessian_approximations import collect_hessians
 
     # collect_hessians writes to partial_run_path (run_path + ".part")
     # We set run_path so partial_run_path points to our desired output location
@@ -301,7 +302,9 @@ def ekfac_results_path(
     cfg.run_path = base_run_path
     cfg.partial_run_path.mkdir(parents=True, exist_ok=True)
 
-    hessian_cfg = HessianConfig(method="kfac", ev_correction=True, use_dataset_labels=True)
+    hessian_cfg = HessianConfig(
+        method="kfac", ev_correction=True, use_dataset_labels=True
+    )
 
     # Phase 1: Covariance collection using collect_hessians
     collect_hessians(

tests/ekfac_tests/test_batch_size_invariance.py

@@ -84,8 +84,12 @@ def compute_traces(batches: list[list[int]]) -> tuple[float, float]:
             computer.run_with_collector_hooks()
 
             # Load covariances
-            A = load_sharded_covariances(index_cfg.partial_run_path / "activation_sharded")
-            G = load_sharded_covariances(index_cfg.partial_run_path / "gradient_sharded")
+            A = load_sharded_covariances(
+                index_cfg.partial_run_path / "activation_sharded"
+            )
+            G = load_sharded_covariances(
+                index_cfg.partial_run_path / "gradient_sharded"
+            )
             n = torch.load(index_cfg.partial_run_path / "total_processed.pt").item()
 
             return (

tests/ekfac_tests/test_compute_ekfac.py

@@ -13,9 +13,7 @@ def test_total_processed_examples(
     total_processed_ground_truth_path = os.path.join(
         ground_truth_covariances_path, "stats.json"
     )
-    total_processed_run_path = os.path.join(
-        ekfac_results_path, "total_processed.pt"
-    )
+    total_processed_run_path = os.path.join(ekfac_results_path, "total_processed.pt")
 
     with open(total_processed_ground_truth_path, "r") as f:
         ground_truth_data = json.load(f)

tests/ekfac_tests/test_covariance.py

@@ -47,8 +47,9 @@ def test_covariances(
     if all_match:
         print(f"{covariance_type} covariances match within tolerance (rtol={rtol})")
     else:
-        error_msg = f"{covariance_type} covariances do not match (rtol={rtol})!\n" + "\n".join(
-            error_details
+        error_msg = (
+            f"{covariance_type} covariances do not match (rtol={rtol})!\n"
+            + "\n".join(error_details)
         )
         assert False, error_msg
 

tests/ekfac_tests/test_eigenvalue_correction.py

@@ -29,9 +29,7 @@ def test_eigenvalue_corrections(
     # load run eigenvalue corrections (sharded)
     lambda_run = load_sharded_covariances(lambda_run_path)
 
-    total_processed_run_path = os.path.join(
-        ekfac_results_path, "total_processed.pt"
-    )
+    total_processed_run_path = os.path.join(ekfac_results_path, "total_processed.pt")
     lambda_device = lambda_run[list(lambda_run.keys())[0]].device
     total = torch.load(total_processed_run_path, map_location=lambda_device)
 
@@ -78,7 +76,9 @@ def test_eigenvalue_corrections(
     if all_match:
         print(f"Eigenvalue corrections match within tolerance (rtol={rtol})")
     elif has_significant_errors:
-        error_msg = f"Eigenvalue corrections do not match (rtol={rtol})!\n" + "\n".join(error_details)
+        error_msg = f"Eigenvalue corrections do not match (rtol={rtol})!\n" + "\n".join(
+            error_details
+        )
         assert False, error_msg
     else:
         print("Eigenvalue corrections: all differences within tolerance")

tests/ekfac_tests/test_fim_accuracy.py

@@ -164,9 +164,15 @@ def test_kfac_fim_accuracy(seq_lengths, num_batches, max_rel_error, sample, tmp_
     computer.forward_backward = fwd_bwd_hessian_factory(index_cfg, hessian_cfg)
     computer.run_with_collector_hooks()
 
-    A_dict_kfac = load_sharded_covariances(index_cfg.partial_run_path / "activation_sharded")
-    G_dict_kfac = load_sharded_covariances(index_cfg.partial_run_path / "gradient_sharded")
-    total_processed_kfac = torch.load(index_cfg.partial_run_path / "total_processed.pt").item()
+    A_dict_kfac = load_sharded_covariances(
+        index_cfg.partial_run_path / "activation_sharded"
+    )
+    G_dict_kfac = load_sharded_covariances(
+        index_cfg.partial_run_path / "gradient_sharded"
+    )
+    total_processed_kfac = torch.load(
+        index_cfg.partial_run_path / "total_processed.pt"
+    ).item()
 
     assert total_processed_kfac == total_processed_exact
 

tests/ekfac_tests/test_utils.py

@@ -10,9 +10,7 @@
 from torch import Tensor
 
 
-def add_tensor_dicts(
-    a: dict[str, Tensor], b: dict[str, Tensor]
-) -> dict[str, Tensor]:
+def add_tensor_dicts(a: dict[str, Tensor], b: dict[str, Tensor]) -> dict[str, Tensor]:
     """Add two dictionaries of tensors element-wise."""
     assert set(a.keys()) == set(b.keys()), "Keys must match"
     return {k: a[k] + b[k] for k in a}