Add paste diff comparison links to TransferLearningAnalysis (facebook#4980)

Summary:
Pull Request resolved: facebook#4980

Follow-up to D92926519. Instead of listing overlapping parameter names inline in the table (which can get long and hard to read), add support for generating a paste diff that shows the current experiment's parameters side-by-side with the source experiment's overlapping parameters.

This follows the same callable-injection pattern used by MetricFetchingErrorsAnalysis: the core analysis accepts an optional `create_diff_paste_callable` that takes (before_content, after_content, title) and returns a diffing URL. When provided, a "Comparison" column is added to the table containing the diff link. The existing "Parameters" column is preserved.

The diff content uses a YAML-like format with `experiment_name` and`parameter_names` keys, making the side-by-side comparison easy to read. The source (old) experiment appears on the left and the current (new) experiment on the right.

Reviewed By: eonofrey

Differential Revision: D95218298

fbshipit-source-id: 372b980372ef5b1f4621011a6e4ad7b96799c015

Author: bernardbeckerman

ax/analysis/healthcheck/tests/test_transfer_learning_analysis.py

@@ -39,6 +39,13 @@ def _make_experiment(
     )
 
 
+def _dummy_create_diff_paste(
+    before_content: str, after_content: str, title: str
+) -> str:
+    """Dummy callable that returns a fake diffing URL."""
+    return "https://www.internalfb.com/intern/diffing/?paste_number=12345"
+
+
 _MOCK_TARGET = "ax.storage.sqa_store.load.identify_transferable_experiments"
 
 
@@ -162,3 +169,81 @@ def test_experiment_name_passed_to_identify(self, mock_identify: object) -> None
         mock_identify.assert_called_once()  # pyre-ignore[16]
         call_kwargs = mock_identify.call_args.kwargs  # pyre-ignore[16]
         self.assertEqual(call_kwargs["experiment_name"], "test_experiment")
+
+    @patch(_MOCK_TARGET)
+    def test_diff_paste_callable_adds_comparison_column(
+        self, mock_identify: object
+    ) -> None:
+        """When create_diff_paste_callable is provided, a 'Comparison' column
+        should be added alongside the existing 'Parameters' column."""
+        experiment = _make_experiment(
+            ["x1", "x2", "x3", "x4"], experiment_type="my_type"
+        )
+        mock_identify.return_value = {  # pyre-ignore[16]
+            "source_exp": TransferLearningMetadata(
+                overlap_parameters=["x1", "x2", "x3"],
+            ),
+        }
+        analysis = TransferLearningAnalysis(
+            create_diff_paste_callable=_dummy_create_diff_paste,
+        )
+        card = analysis.compute(experiment=experiment)
+        self.assertIn("Comparison", card.df.columns)
+        self.assertIn("Parameters", card.df.columns)
+        self.assertIn("diffing", card.df.iloc[0]["Comparison"])
+        self.assertEqual(card.df.iloc[0]["Parameters"], "x1, x2, x3")
+
+    @patch(_MOCK_TARGET)
+    def test_diff_paste_callable_receives_correct_content(
+        self, mock_identify: object
+    ) -> None:
+        """Verify before/after content includes experiment name headers and
+        sorted parameters."""
+        experiment = _make_experiment(
+            ["alpha", "beta", "gamma"], experiment_type="my_type"
+        )
+        mock_identify.return_value = {  # pyre-ignore[16]
+            "source_exp": TransferLearningMetadata(
+                overlap_parameters=["gamma", "alpha"],
+            ),
+        }
+        captured_args: list[tuple[str, str, str]] = []
+
+        def _capture_callable(before: str, after: str, title: str) -> str:
+            captured_args.append((before, after, title))
+            return "https://example.com/diff"
+
+        analysis = TransferLearningAnalysis(
+            create_diff_paste_callable=_capture_callable,
+        )
+        analysis.compute(experiment=experiment)
+        self.assertEqual(len(captured_args), 1)
+        before, after, title = captured_args[0]
+        # Before should have source experiment in YAML format
+        self.assertEqual(
+            before,
+            "experiment_name: source_exp (old)\nparameter_names:\n  - alpha\n  - gamma",
+        )
+        # After should have current experiment in YAML format
+        self.assertEqual(
+            after,
+            "experiment_name: test_experiment (new)\n"
+            "parameter_names:\n  - alpha\n  - beta\n  - gamma",
+        )
+        self.assertIn("source_exp", title)
+        self.assertIn("test_experiment", title)
+
+    @patch(_MOCK_TARGET)
+    def test_no_callable_has_no_comparison_column(self, mock_identify: object) -> None:
+        """Without callable, the 'Parameters' column should be present
+        but no 'Comparison' column."""
+        experiment = _make_experiment(["x1", "x2", "x3"], experiment_type="my_type")
+        mock_identify.return_value = {  # pyre-ignore[16]
+            "source_exp": TransferLearningMetadata(
+                overlap_parameters=["x1", "x2"],
+            ),
+        }
+        analysis = TransferLearningAnalysis()
+        card = analysis.compute(experiment=experiment)
+        self.assertIn("Parameters", card.df.columns)
+        self.assertNotIn("Comparison", card.df.columns)

ax/analysis/healthcheck/transfer_learning_analysis.py

@@ -8,7 +8,8 @@
 from __future__ import annotations
 
 import json
-from typing import final, TYPE_CHECKING
+from collections.abc import Callable
+from typing import Any, final, TYPE_CHECKING
 
 import markdown as md
 import pandas as pd
@@ -46,11 +47,26 @@ def __init__(
         overlap_threshold: float = 0.50,
         max_num_exps: int = 10,
         config: SQAConfig | None = None,
+        create_diff_paste_callable: Callable[[str, str, str], str] | None = None,
     ) -> None:
+        """
+        Args:
+            experiment_types: List of experiment types to search for.
+            overlap_threshold: Minimum proportion of overlapping parameters.
+            max_num_exps: Max number of transferable experiments to return.
+            config: SQAConfig to use for the query.
+            create_diff_paste_callable: A function that takes
+                (before_content, after_content, title) and returns a URL to a
+                paste diff comparing the current experiment's parameters with
+                a source experiment's overlapping parameters. If provided, a
+                "Comparison" column is added to the table containing the diff
+                link.
+        """
         self.experiment_types = experiment_types
         self.overlap_threshold = overlap_threshold
         self.max_num_exps = max_num_exps
         self.config = config
+        self.create_diff_paste_callable = create_diff_paste_callable
 
     @override
     def compute(
@@ -104,23 +120,42 @@ def compute(
             )
 
         total_parameters = len(experiment.search_space.parameters)
+        current_params_sorted = sorted(experiment.search_space.parameters.keys())
 
-        rows = []
+        rows: list[dict[str, Any]] = []
         for exp_name, metadata in transferable_experiments.items():
             overlap_count = len(metadata.overlap_parameters)
             overlap_pct = (
                 (overlap_count / total_parameters * 100)
                 if total_parameters > 0
                 else 0.0
             )
-            rows.append(
-                {
-                    "Experiment": exp_name,
-                    "Overlapping Parameters": overlap_count,
-                    "Overlap (%)": round(overlap_pct, 1),
-                    "Parameters": ", ".join(sorted(metadata.overlap_parameters)),
-                }
-            )
+            overlap_sorted = sorted(metadata.overlap_parameters)
+            row: dict[str, Any] = {
+                "Experiment": exp_name,
+                "Overlapping Parameters": overlap_count,
+                "Overlap (%)": round(overlap_pct, 1),
+                "Parameters": ", ".join(overlap_sorted),
+            }
+
+            if self.create_diff_paste_callable is not None:
+                create_diff_paste = self.create_diff_paste_callable
+                before_params = "\n".join(f"  - {p}" for p in overlap_sorted)
+                before_content = (
+                    f"experiment_name: {exp_name} (old)\n"
+                    f"parameter_names:\n{before_params}"
+                )
+                after_params = "\n".join(f"  - {p}" for p in current_params_sorted)
+                after_content = (
+                    f"experiment_name: {experiment.name} (new)\n"
+                    f"parameter_names:\n{after_params}"
+                )
+                title = f"Parameter comparison: {experiment.name} vs {exp_name}"
+                row["Comparison"] = create_diff_paste(
+                    before_content, after_content, title
+                )
+
+            rows.append(row)
 
         df = pd.DataFrame(rows)
 

ax/analysis/overview.py

@@ -5,6 +5,7 @@
 
 # pyre-strict
 
+from collections.abc import Callable
 from typing import Any, final
 
 from ax.adapter.base import Adapter
@@ -116,6 +117,7 @@ def __init__(
         tier_metadata: dict[str, Any] | None = None,
         model_fit_threshold: float | None = None,
         sqa_config: Any = None,
+        create_diff_paste_callable: Callable[[str, str, str], str] | None = None,
     ) -> None:
         super().__init__()
         self.can_generate = can_generate
@@ -127,6 +129,7 @@ def __init__(
         self.tier_metadata = tier_metadata
         self.model_fit_threshold = model_fit_threshold
         self.sqa_config = sqa_config
+        self.create_diff_paste_callable = create_diff_paste_callable
 
     @override
     def validate_applicable_state(
@@ -232,7 +235,10 @@ def compute(
             if not has_batch_trials
             else None,
             BaselineImprovementAnalysis() if not has_batch_trials else None,
-            TransferLearningAnalysis(config=self.sqa_config),
+            TransferLearningAnalysis(
+                config=self.sqa_config,
+                create_diff_paste_callable=self.create_diff_paste_callable,
+            ),
             *[
                 SearchSpaceAnalysis(trial_index=trial.index)
                 for trial in candidate_trials