No public description

fehiepsi · copybara-github · commit befb6b9fdd06 · 2025-03-31T08:29:09.000-07:00
PiperOrigin-RevId: 742274292
Change-Id: Ifbaefaf361142577e03bca03d0d1d197d7f140fd
diff --git a/lightweight_mmm/core/baseline/intercept.py b/lightweight_mmm/core/baseline/intercept.py
@@ -52,4 +52,4 @@ def simple_intercept(
         fn=custom_priors.get(priors.INTERCEPT,
                              default_priors[priors.INTERCEPT]),
     )
-  return intercept
+  return jnp.asarray(intercept)
diff --git a/lightweight_mmm/models.py b/lightweight_mmm/models.py
@@ -420,8 +420,9 @@ def media_mix_model(
       plate_prefixes = ("extra_feature", "geo")
       extra_features_einsum = "tfg, fg -> tg"  # t = time, f = feature, g = geo
       extra_features_plates_shape = (extra_features.shape[1], *geo_shape)
-    with numpyro.plate_stack(plate_prefixes,
-                             sizes=extra_features_plates_shape):
+    with numpyro.plate_stack(
+        str(plate_prefixes), sizes=list(extra_features_plates_shape)
+    ):
       coef_extra_features = numpyro.sample(
           name=_COEF_EXTRA_FEATURES,
           fn=custom_priors.get(

Original file line number	Diff line number	Diff line change
`@@ -52,4 +52,4 @@ def simple_intercept(`
`52`	`52`	`fn=custom_priors.get(priors.INTERCEPT,`
`53`	`53`	`default_priors[priors.INTERCEPT]),`
`54`	`54`	`)`
`55`		`- return intercept`
	`55`	`+ return jnp.asarray(intercept)`