google-meridian 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

meridian/model/prior_sampler.py

@@ -48,15 +48,14 @@ def _get_tau_g(
     parameter with zero at position `baseline_geo_idx` and matching
     `tau_g_excl_baseline` elsewhere.
   """
-  rank = len(tau_g_excl_baseline.shape)
-  shape = tau_g_excl_baseline.shape[:-1] + [1] if rank != 1 else 1
+  shape = tau_g_excl_baseline.shape[:-1] + (1,)
   tau_g = backend.concatenate(
       [
           tau_g_excl_baseline[..., :baseline_geo_idx],
           backend.zeros(shape, dtype=tau_g_excl_baseline.dtype),
           tau_g_excl_baseline[..., baseline_geo_idx:],
       ],
-      axis=rank - 1,
+      axis=-1,
   )
   return backend.tfd.Deterministic(tau_g, name="tau_g")
 
@@ -70,15 +69,13 @@ class PriorDistributionSampler:
   def _sample_media_priors(
       self,
       n_draws: int,
-      seed: int | None = None,
+      rng_handler: backend.RNGHandler,
   ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the media variables.
 
     Args:
       n_draws: Number of samples drawn from the prior distribution.
-      seed: Used to set the seed for reproducible results. For more information,
-        see [PRNGS and seeds]
-        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      rng_handler: The backend-agnostic RNG handler managing the seed state.
 
     Returns:
       A mapping of media parameter names to a tensor of shape `[n_draws, n_geos,
@@ -89,31 +86,47 @@ class PriorDistributionSampler:
 
     prior = mmm.prior_broadcast
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
+
     media_vars = {
-        constants.ALPHA_M: prior.alpha_m.sample(**sample_kwargs),
-        constants.EC_M: prior.ec_m.sample(**sample_kwargs),
-        constants.ETA_M: prior.eta_m.sample(**sample_kwargs),
-        constants.SLOPE_M: prior.slope_m.sample(**sample_kwargs),
+        constants.ALPHA_M: prior.alpha_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.EC_M: prior.ec_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.ETA_M: prior.eta_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.SLOPE_M: prior.slope_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
     }
     beta_gm_dev = backend.tfd.Sample(
         backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_media_channels],
         name=constants.BETA_GM_DEV,
-    ).sample(**sample_kwargs)
+    ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
 
     prior_type = mmm.model_spec.effective_media_prior_type
     if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
-      media_vars[constants.BETA_M] = prior.beta_m.sample(**sample_kwargs)
+      media_vars[constants.BETA_M] = prior.beta_m.sample(
+          sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+      )
     else:
       if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
-        treatment_parameter_m = prior.roi_m.sample(**sample_kwargs)
+        treatment_parameter_m = prior.roi_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         media_vars[constants.ROI_M] = treatment_parameter_m
       elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
-        treatment_parameter_m = prior.mroi_m.sample(**sample_kwargs)
+        treatment_parameter_m = prior.mroi_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         media_vars[constants.MROI_M] = treatment_parameter_m
       elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
-        treatment_parameter_m = prior.contribution_m.sample(**sample_kwargs)
+        treatment_parameter_m = prior.contribution_m.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         media_vars[constants.CONTRIBUTION_M] = treatment_parameter_m
       else:
         raise ValueError(f"Unsupported prior type: {prior_type}")
@@ -125,7 +138,7 @@ class PriorDistributionSampler:
           alpha=media_vars[constants.ALPHA_M],
           ec=media_vars[constants.EC_M],
           slope=media_vars[constants.SLOPE_M],
-          decay_functions=mmm.adstock_decay_spec.media
+          decay_functions=mmm.adstock_decay_spec.media,
       )
       linear_predictor_counterfactual_difference = (
           mmm.linear_predictor_counterfactual_difference_media(
@@ -144,7 +157,7 @@ class PriorDistributionSampler:
       )
       media_vars[constants.BETA_M] = backend.tfd.Deterministic(
           beta_m_value, name=constants.BETA_M
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
 
     beta_eta_combined = (
         media_vars[constants.BETA_M][..., backend.newaxis, :]
@@ -157,22 +170,20 @@ class PriorDistributionSampler:
     )
     media_vars[constants.BETA_GM] = backend.tfd.Deterministic(
         beta_gm_value, name=constants.BETA_GM
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
 
     return media_vars
 
   def _sample_rf_priors(
       self,
       n_draws: int,
-      seed: int | None = None,
+      rng_handler: backend.RNGHandler,
   ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the RF variables.
 
     Args:
       n_draws: Number of samples drawn from the prior distribution.
-      seed: Used to set the seed for reproducible results. For more information,
-        see [PRNGS and seeds]
-        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      rng_handler: The backend-agnostic RNG handler managing the seed state.
 
     Returns:
       A mapping of RF parameter names to a tensor of shape
@@ -183,31 +194,47 @@ class PriorDistributionSampler:
 
     prior = mmm.prior_broadcast
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
+
     rf_vars = {
-        constants.ALPHA_RF: prior.alpha_rf.sample(**sample_kwargs),
-        constants.EC_RF: prior.ec_rf.sample(**sample_kwargs),
-        constants.ETA_RF: prior.eta_rf.sample(**sample_kwargs),
-        constants.SLOPE_RF: prior.slope_rf.sample(**sample_kwargs),
+        constants.ALPHA_RF: prior.alpha_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.EC_RF: prior.ec_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.ETA_RF: prior.eta_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.SLOPE_RF: prior.slope_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
     }
     beta_grf_dev = backend.tfd.Sample(
         backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_rf_channels],
         name=constants.BETA_GRF_DEV,
-    ).sample(**sample_kwargs)
+    ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
 
     prior_type = mmm.model_spec.effective_rf_prior_type
     if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
-      rf_vars[constants.BETA_RF] = prior.beta_rf.sample(**sample_kwargs)
+      rf_vars[constants.BETA_RF] = prior.beta_rf.sample(
+          sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+      )
     else:
       if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
-        treatment_parameter_rf = prior.roi_rf.sample(**sample_kwargs)
+        treatment_parameter_rf = prior.roi_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         rf_vars[constants.ROI_RF] = treatment_parameter_rf
       elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
-        treatment_parameter_rf = prior.mroi_rf.sample(**sample_kwargs)
+        treatment_parameter_rf = prior.mroi_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         rf_vars[constants.MROI_RF] = treatment_parameter_rf
       elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
-        treatment_parameter_rf = prior.contribution_rf.sample(**sample_kwargs)
+        treatment_parameter_rf = prior.contribution_rf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        )
         rf_vars[constants.CONTRIBUTION_RF] = treatment_parameter_rf
       else:
         raise ValueError(f"Unsupported prior type: {prior_type}")
@@ -240,7 +267,7 @@ class PriorDistributionSampler:
       rf_vars[constants.BETA_RF] = backend.tfd.Deterministic(
           beta_rf_value,
           name=constants.BETA_RF,
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
 
     beta_eta_combined = (
         rf_vars[constants.BETA_RF][..., backend.newaxis, :]
@@ -253,22 +280,20 @@ class PriorDistributionSampler:
     )
     rf_vars[constants.BETA_GRF] = backend.tfd.Deterministic(
         beta_grf_value, name=constants.BETA_GRF
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
 
     return rf_vars
 
   def _sample_organic_media_priors(
       self,
       n_draws: int,
-      seed: int | None = None,
+      rng_handler: backend.RNGHandler,
   ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of organic media variables.
 
     Args:
       n_draws: Number of samples drawn from the prior distribution.
-      seed: Used to set the seed for reproducible results. For more information,
-        see [PRNGS and seeds]
-        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      rng_handler: The backend-agnostic RNG handler managing the seed state.
 
     Returns:
       A mapping of organic media parameter names to a tensor of shape
@@ -279,27 +304,37 @@ class PriorDistributionSampler:
 
     prior = mmm.prior_broadcast
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
+
     organic_media_vars = {
-        constants.ALPHA_OM: prior.alpha_om.sample(**sample_kwargs),
-        constants.EC_OM: prior.ec_om.sample(**sample_kwargs),
-        constants.ETA_OM: prior.eta_om.sample(**sample_kwargs),
-        constants.SLOPE_OM: prior.slope_om.sample(**sample_kwargs),
+        constants.ALPHA_OM: prior.alpha_om.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.EC_OM: prior.ec_om.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.ETA_OM: prior.eta_om.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.SLOPE_OM: prior.slope_om.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
     }
     beta_gom_dev = backend.tfd.Sample(
         backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_organic_media_channels],
         name=constants.BETA_GOM_DEV,
-    ).sample(**sample_kwargs)
+    ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
 
     prior_type = mmm.model_spec.organic_media_prior_type
     if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
       organic_media_vars[constants.BETA_OM] = prior.beta_om.sample(
-          **sample_kwargs
+          sample_shape=sample_shape, seed=rng_handler.get_next_seed()
       )
     elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
       organic_media_vars[constants.CONTRIBUTION_OM] = (
-          prior.contribution_om.sample(**sample_kwargs)
+          prior.contribution_om.sample(
+              sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+          )
       )
       incremental_outcome_om = (
           organic_media_vars[constants.CONTRIBUTION_OM] * mmm.total_outcome
@@ -321,7 +356,7 @@ class PriorDistributionSampler:
       organic_media_vars[constants.BETA_OM] = backend.tfd.Deterministic(
           beta_om_value,
           name=constants.BETA_OM,
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
     else:
       raise ValueError(f"Unsupported prior type: {prior_type}")
 
@@ -337,22 +372,20 @@ class PriorDistributionSampler:
     )
     organic_media_vars[constants.BETA_GOM] = backend.tfd.Deterministic(
         beta_gom_value, name=constants.BETA_GOM
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
 
     return organic_media_vars
 
   def _sample_organic_rf_priors(
       self,
       n_draws: int,
-      seed: int | None = None,
+      rng_handler: backend.RNGHandler,
   ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the organic RF variables.
 
     Args:
       n_draws: Number of samples drawn from the prior distribution.
-      seed: Used to set the seed for reproducible results. For more information,
-        see [PRNGS and seeds]
-        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      rng_handler: The backend-agnostic RNG handler managing the seed state.
 
     Returns:
       A mapping of organic RF parameter names to a tensor of shape
@@ -363,27 +396,37 @@ class PriorDistributionSampler:
 
     prior = mmm.prior_broadcast
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
+
     organic_rf_vars = {
-        constants.ALPHA_ORF: prior.alpha_orf.sample(**sample_kwargs),
-        constants.EC_ORF: prior.ec_orf.sample(**sample_kwargs),
-        constants.ETA_ORF: prior.eta_orf.sample(**sample_kwargs),
-        constants.SLOPE_ORF: prior.slope_orf.sample(**sample_kwargs),
+        constants.ALPHA_ORF: prior.alpha_orf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.EC_ORF: prior.ec_orf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.ETA_ORF: prior.eta_orf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.SLOPE_ORF: prior.slope_orf.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
     }
     beta_gorf_dev = backend.tfd.Sample(
         backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_organic_rf_channels],
         name=constants.BETA_GORF_DEV,
-    ).sample(**sample_kwargs)
+    ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
 
     prior_type = mmm.model_spec.organic_media_prior_type
     if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
       organic_rf_vars[constants.BETA_ORF] = prior.beta_orf.sample(
-          **sample_kwargs
+          sample_shape=sample_shape, seed=rng_handler.get_next_seed()
       )
     elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
       organic_rf_vars[constants.CONTRIBUTION_ORF] = (
-          prior.contribution_orf.sample(**sample_kwargs)
+          prior.contribution_orf.sample(
+              sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+          )
       )
       incremental_outcome_orf = (
           organic_rf_vars[constants.CONTRIBUTION_ORF] * mmm.total_outcome
@@ -406,7 +449,7 @@ class PriorDistributionSampler:
       organic_rf_vars[constants.BETA_ORF] = backend.tfd.Deterministic(
           beta_orf_value,
           name=constants.BETA_ORF,
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
     else:
       raise ValueError(f"Unsupported prior type: {prior_type}")
 
@@ -422,22 +465,20 @@ class PriorDistributionSampler:
     )
     organic_rf_vars[constants.BETA_GORF] = backend.tfd.Deterministic(
         beta_gorf_value, name=constants.BETA_GORF
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
 
     return organic_rf_vars
 
   def _sample_non_media_treatments_priors(
       self,
       n_draws: int,
-      seed: int | None = None,
+      rng_handler: backend.RNGHandler,
   ) -> Mapping[str, backend.Tensor]:
     """Draws from the prior distributions of the non-media treatment variables.
 
     Args:
       n_draws: Number of samples drawn from the prior distribution.
-      seed: Used to set the seed for reproducible results. For more information,
-        see [PRNGS and seeds]
-        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      rng_handler: The backend-agnostic RNG handler managing the seed state.
 
     Returns:
       A mapping of non-media treatment parameter names to a tensor of shape
@@ -448,23 +489,27 @@ class PriorDistributionSampler:
 
     prior = mmm.prior_broadcast
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
+
     non_media_treatments_vars = {
-        constants.XI_N: prior.xi_n.sample(**sample_kwargs),
+        constants.XI_N: prior.xi_n.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
     }
     gamma_gn_dev = backend.tfd.Sample(
         backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_non_media_channels],
         name=constants.GAMMA_GN_DEV,
-    ).sample(**sample_kwargs)
+    ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
     prior_type = mmm.model_spec.non_media_treatments_prior_type
     if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
       non_media_treatments_vars[constants.GAMMA_N] = prior.gamma_n.sample(
-          **sample_kwargs
+          sample_shape=sample_shape, seed=rng_handler.get_next_seed()
       )
     elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
       non_media_treatments_vars[constants.CONTRIBUTION_N] = (
-          prior.contribution_n.sample(**sample_kwargs)
+          prior.contribution_n.sample(
+              sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+          )
       )
       incremental_outcome_n = (
           non_media_treatments_vars[constants.CONTRIBUTION_N]
@@ -485,7 +530,7 @@ class PriorDistributionSampler:
       )
       non_media_treatments_vars[constants.GAMMA_N] = backend.tfd.Deterministic(
           gamma_n_value, name=constants.GAMMA_N
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
     else:
       raise ValueError(f"Unsupported prior type: {prior_type}")
     non_media_treatments_vars[constants.GAMMA_GN] = backend.tfd.Deterministic(
@@ -493,7 +538,7 @@ class PriorDistributionSampler:
         + non_media_treatments_vars[constants.XI_N][..., backend.newaxis, :]
         * gamma_gn_dev,
         name=constants.GAMMA_GN,
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
     return non_media_treatments_vars
 
   def _sample_prior(
@@ -509,21 +554,28 @@ class PriorDistributionSampler:
     if seed is not None:
       backend.set_random_seed(seed)
 
+    rng_handler = backend.RNGHandler(seed)
+
     prior = mmm.prior_broadcast
     # `sample_shape` is prepended to the shape of each BatchBroadcast in `prior`
     # when it is sampled.
     sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
 
-    tau_g_excl_baseline = prior.tau_g_excl_baseline.sample(**sample_kwargs)
+    tau_g_excl_baseline = prior.tau_g_excl_baseline.sample(
+        sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+    )
     base_vars = {
-        constants.KNOT_VALUES: prior.knot_values.sample(**sample_kwargs),
-        constants.SIGMA: prior.sigma.sample(**sample_kwargs),
+        constants.KNOT_VALUES: prior.knot_values.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
+        constants.SIGMA: prior.sigma.sample(
+            sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+        ),
         constants.TAU_G: (
             _get_tau_g(
                 tau_g_excl_baseline=tau_g_excl_baseline,
                 baseline_geo_idx=mmm.baseline_geo_idx,
-            ).sample()
+            ).sample(seed=rng_handler.get_next_seed())
         ),
     }
 
@@ -534,49 +586,53 @@ class PriorDistributionSampler:
             backend.to_tensor(mmm.knot_info.weights),
         ),
         name=constants.MU_T,
-    ).sample()
+    ).sample(seed=rng_handler.get_next_seed())
 
     # Omit gamma_c, xi_c, and gamma_gc parameters from sampled distributions if
     # there are no control variables in the model.
     if mmm.n_controls:
       base_vars |= {
-          constants.GAMMA_C: prior.gamma_c.sample(**sample_kwargs),
-          constants.XI_C: prior.xi_c.sample(**sample_kwargs),
+          constants.GAMMA_C: prior.gamma_c.sample(
+              sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+          ),
+          constants.XI_C: prior.xi_c.sample(
+              sample_shape=sample_shape, seed=rng_handler.get_next_seed()
+          ),
       }
 
       gamma_gc_dev = backend.tfd.Sample(
           backend.tfd.Normal(0, 1),
           [mmm.n_geos, mmm.n_controls],
           name=constants.GAMMA_GC_DEV,
-      ).sample(**sample_kwargs)
+      ).sample(sample_shape=sample_shape, seed=rng_handler.get_next_seed())
       base_vars[constants.GAMMA_GC] = backend.tfd.Deterministic(
           base_vars[constants.GAMMA_C][..., backend.newaxis, :]
           + base_vars[constants.XI_C][..., backend.newaxis, :] * gamma_gc_dev,
           name=constants.GAMMA_GC,
-      ).sample()
+      ).sample(seed=rng_handler.get_next_seed())
 
     media_vars = (
-        self._sample_media_priors(n_draws, seed)
+        self._sample_media_priors(n_draws, rng_handler)
         if mmm.media_tensors.media is not None
         else {}
     )
     rf_vars = (
-        self._sample_rf_priors(n_draws, seed)
+        self._sample_rf_priors(n_draws, rng_handler)
         if mmm.rf_tensors.reach is not None
         else {}
     )
     organic_media_vars = (
-        self._sample_organic_media_priors(n_draws, seed)
+        self._sample_organic_media_priors(n_draws, rng_handler)
         if mmm.organic_media_tensors.organic_media is not None
         else {}
     )
     organic_rf_vars = (
-        self._sample_organic_rf_priors(n_draws, seed)
+        self._sample_organic_rf_priors(n_draws, rng_handler)
         if mmm.organic_rf_tensors.organic_reach is not None
         else {}
     )
     non_media_treatments_vars = (
-        self._sample_non_media_treatments_priors(n_draws, seed)
+        self._sample_non_media_treatments_priors(n_draws, rng_handler)
         if mmm.non_media_treatments_normalized is not None
         else {}
     )

meridian/model/spec.py

@@ -203,14 +203,13 @@ class ModelSpec:
       non-media value will be scaled by population. If `None`, then no non-media
       variables are scaled by population. Default: `None`.
     adstock_decay_spec: A string or mapping specifying the adstock decay
-      function for each media, RF, organic media and organic RF channel.
-      * If a string, must be either `'geometric'` or `'binomial'`, specifying
-        that decay function for all channels.
-      * If a mapping, keys should be channel names and values should be
-        `'geometric'` or `'binomial'`, with each key-value pair denoting the
-        adstock decay function to use for that channel. Channels that are not
-        specified in the mapping default to using 'geometric'.
-      Default: `'geometric'`.
+      function for each media, RF, organic media and organic RF channel. If a
+      string, must be either `'geometric'` or `'binomial'`, specifying that
+      decay function for all channels. If a mapping, keys should be channel
+      names and values should be `'geometric'` or `'binomial'`, with each
+      key-value pair denoting the adstock decay function to use for that
+      channel. Channels that are not specified in the mapping default to using
+      'geometric'. Default: `'geometric'`.
     enable_aks: A boolean indicating whether to use the Automatic Knot Selection
       algorithm to select an optimal number of knots for running the model
       instead of the default 1 for national models and n_times for geo models.

meridian/model/transformers.py

@@ -196,11 +196,19 @@ class KpiTransformer(TensorTransformer):
         each geo, used to to compute the population scale factors.
     """
     self._population = population
-    population_scaled_kpi = backend.divide_no_nan(
+    self._population_scaled_kpi = backend.divide_no_nan(
         kpi, self._population[:, backend.newaxis]
     )
-    self._population_scaled_mean = backend.reduce_mean(population_scaled_kpi)
-    self._population_scaled_stdev = backend.reduce_std(population_scaled_kpi)
+    self._population_scaled_mean = backend.reduce_mean(
+        self._population_scaled_kpi
+    )
+    self._population_scaled_stdev = backend.reduce_std(
+        self._population_scaled_kpi
+    )
+
+  @property
+  def population_scaled_kpi(self):
+    return self._population_scaled_kpi
 
   @property
   def population_scaled_mean(self):

meridian/version.py

@@ -14,4 +14,4 @@
 
 """Module for Meridian version."""
 
-__version__ = "1.2.1"
+__version__ = "1.3.1"

schema/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module containing MMM schema library."""
+
+from schema import serde
+from schema import utils

schema/serde/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A serialization and deserialization library for Meridian models.
+
+For entry points API, see `meridian_serde` module docs.
+"""
+
+from schema.serde import constants
+from schema.serde import distribution
+from schema.serde import eda_spec
+from schema.serde import hyperparameters
+from schema.serde import inference_data
+from schema.serde import meridian_serde
+from schema.serde import serde