Internal change

PiperOrigin-RevId: 544062601
Change-Id: Ifbc8967d1d78fbce458700ec6fd9f562c977250c

lightweight_mmm/lightweight_mmm.py

@@ -374,7 +374,7 @@ def fit(
 
     self.custom_priors = custom_priors
     if media_names is not None:
-      self.media_names = media_names
+      self.media_names = list(media_names)
     else:
       self.media_names = [f"channel_{i}" for i in range(media.shape[1])]
     self.n_media_channels = media.shape[1]

lightweight_mmm/lightweight_mmm_test.py

@@ -21,6 +21,7 @@
 import jax.numpy as jnp
 import numpy as np
 import numpyro.distributions as dist
+import pandas as pd
 
 from lightweight_mmm import lightweight_mmm
 from lightweight_mmm import models
@@ -463,5 +464,19 @@ def test_fitted_mmm_does_not_equal_default_mmm(self, media_mix_model):
     fitted_mmm_object = getattr(self, media_mix_model)
     self.assertNotEqual(default_mmm_object, fitted_mmm_object)
 
+  def test_equality_function_works_with_media_names_as_pandas_index(self):
+    mmm_object = lightweight_mmm.LightweightMMM()
+    mmm_object.fit(
+        media=jnp.ones((50, 5)),
+        target=jnp.ones(50),
+        media_prior=jnp.ones(5) * 50,
+        extra_features=jnp.ones((50, 2)),
+        number_warmup=2,
+        number_samples=4,
+        number_chains=1,
+        media_names=pd.Index([f'channel_{i}' for i in range(5)]))
+
+    self.assertEqual(mmm_object, mmm_object)
+
 if __name__ == "__main__":
   absltest.main()