diff --git a/meridian/model/eda/meridian_eda.py b/meridian/model/eda/meridian_eda.py
index c985f9289..eee6d50b9 100644
--- a/meridian/model/eda/meridian_eda.py
+++ b/meridian/model/eda/meridian_eda.py
@@ -35,7 +35,6 @@
 import pandas as pd
 import xarray as xr
 
-
 if TYPE_CHECKING:
   from meridian.model import model  # pylint: disable=g-bad-import-order,g-import-not-at-top
 
@@ -2274,10 +2273,13 @@ def _get_plot_data_for_heatmap(
   # Maintain the original order of variables.
   plot_variables = [v for v in all_variables if v in selected_vars]
 
-  return df[
-      df[eda_constants.VARIABLE_1].isin(plot_variables)
-      & df[eda_constants.VARIABLE_2].isin(plot_variables)
-  ], plot_variables
+  return (
+      df[
+          df[eda_constants.VARIABLE_1].isin(plot_variables)
+          & df[eda_constants.VARIABLE_2].isin(plot_variables)
+      ],
+      plot_variables,
+  )
 
 
 def _format_explanation_for_html(explanation: str) -> str:
diff --git a/meridian/model/eda/meridian_eda_test.py b/meridian/model/eda/meridian_eda_test.py
index a0f5bd131..b526dc5a9 100644
--- a/meridian/model/eda/meridian_eda_test.py
+++ b/meridian/model/eda/meridian_eda_test.py
@@ -34,6 +34,7 @@
 from meridian.model.eda import eda_spec
 from meridian.model.eda import meridian_eda
 from meridian.model.eda import sampling_eda_engine
+
 import numpy as np
 import pandas as pd
 import xarray as xr