Merge pull request #17 from ryusudol/dev

Dev

Author: ryusudol

.github/workflows/ci.yaml

@@ -2,7 +2,7 @@ name: CI
 
 on:
   push:
-    branches: [main, dev]
+    branches: [main]
     paths:
       - "cka/**"
       - "tests/**"

.gitignore

@@ -57,3 +57,5 @@ cka_checkpoint.pt
 .claude/
 CLAUDE.md
 uv.lock
+
+*.pdf

cka/cka.py

@@ -242,7 +242,7 @@ def compare(
         dataloader: DataLoader,
         dataloader2: DataLoader | None = None,
         progress: bool = True,
-        callback: Callable[[int, int, torch.Tensor], None] | None = None,
+        verbose: bool = False,
     ) -> torch.Tensor:
         if not self._hook_handles:
             raise RuntimeError("Hooks not registered.")
@@ -296,9 +296,11 @@ def compare(
 
                 self._accumulate_hsic(hsic_xy, hsic_xx, hsic_yy)
 
-                if callback is not None:
+                if verbose:
                     current_cka = self._compute_cka_matrix(hsic_xy, hsic_xx, hsic_yy)
-                    callback(batch_idx, total_batches, current_cka)
+                    print(
+                        f"Batch {batch_idx + 1}/{total_batches} - Mean CKA: {current_cka.mean().item():.4f}"
+                    )
 
         return self._compute_cka_matrix(hsic_xy, hsic_xx, hsic_yy)
 

examples/performance_comparison.ipynb …mples/computation_speed_comparison.ipynbexamples/performance_comparison.ipynb renamed to examples/computation_speed_comparison.ipynb examples/performance_comparison.ipynb renamed to examples/computation_speed_comparison.ipynb

@@ -336,19 +336,16 @@ def test_progress_true(self, model1, model2, dataloader):
 
         assert result.shape == (3, 3)
 
-    def test_callback_called(self, model1, model2, dataloader):
+    def test_verbose_output(self, model1, model2, dataloader, capsys):
         cka = CKA(model1, model2)
-        callback_calls = []
 
-        def callback(batch_idx, total_batches, cka_matrix):
-            callback_calls.append((batch_idx, total_batches, cka_matrix.clone()))
+        result = cka(dataloader, verbose=True, progress=False)
 
-        result = cka(dataloader, callback=callback)
-
-        assert len(callback_calls) == 4
-        assert callback_calls[0][0] == 0
-        assert callback_calls[0][1] == 4
-        assert callback_calls[-1][2].shape == result.shape
+        captured = capsys.readouterr()
+        assert "Batch 1/4" in captured.out
+        assert "Batch 4/4" in captured.out
+        assert "Mean CKA:" in captured.out
+        assert result.shape[0] > 0
 
     def test_two_dataloaders(self, model1, model2, dataloader):
         x2 = torch.randn(32, 10)