google-meridian 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

meridian/model/model.py

@@ -27,12 +27,13 @@ from meridian.data import time_coordinates as tc
 from meridian.model import adstock_hill
 from meridian.model import knots
 from meridian.model import media
+from meridian.model import posterior_sampler
 from meridian.model import prior_distribution
+from meridian.model import prior_sampler
 from meridian.model import spec
 from meridian.model import transformers
 import numpy as np
 import tensorflow as tf
-import tensorflow_probability as tfp
 
 
 __all__ = [
@@ -49,12 +50,8 @@ class NotFittedModelError(Exception):
   """Model has not been fitted."""
 
 
-class MCMCSamplingError(Exception):
-  """The Markov Chain Monte Carlo (MCMC) sampling failed."""
-
-
-class MCMCOOMError(Exception):
-  """The Markov Chain Monte Carlo (MCMC) exceeds memory limits."""
+MCMCSamplingError = posterior_sampler.MCMCSamplingError
+MCMCOOMError = posterior_sampler.MCMCOOMError
 
 
 def _warn_setting_national_args(**kwargs):
@@ -70,43 +67,6 @@ def _warn_setting_national_args(**kwargs):
       )
 
 
-def _get_tau_g(
-    tau_g_excl_baseline: tf.Tensor, baseline_geo_idx: int
-) -> tfp.distributions.Distribution:
-  """Computes `tau_g` from `tau_g_excl_baseline`.
-
-  This function computes `tau_g` by inserting a column of zeros at the
-  `baseline_geo` position in `tau_g_excl_baseline`.
-
-  Args:
-    tau_g_excl_baseline: A tensor of shape `[..., n_geos - 1]` for the
-      user-defined dimensions of the `tau_g` parameter distribution.
-    baseline_geo_idx: The index of the baseline geo to be set to zero.
-
-  Returns:
-    A tensor of shape `[..., n_geos]` with the final distribution of the `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
-  tau_g = tf.concat(
-      [
-          tau_g_excl_baseline[..., :baseline_geo_idx],
-          tf.zeros(shape, dtype=tau_g_excl_baseline.dtype),
-          tau_g_excl_baseline[..., baseline_geo_idx:],
-      ],
-      axis=rank - 1,
-  )
-  return tfp.distributions.Deterministic(tau_g, name="tau_g")
-
-
-@tf.function(autograph=False, jit_compile=True)
-def _xla_windowed_adaptive_nuts(**kwargs):
-  """XLA wrapper for windowed_adaptive_nuts."""
-  return tfp.experimental.mcmc.windowed_adaptive_nuts(**kwargs)
-
-
 class Meridian:
   """Contains the main functionality for fitting the Meridian MMM model.
 
@@ -452,6 +412,18 @@ class Meridian:
         total_spend=agg_total_spend,
     )
 
+  @functools.cached_property
+  def prior_sampler_callable(self) -> prior_sampler.PriorDistributionSampler:
+    """A `PriorDistributionSampler` callable bound to this model."""
+    return prior_sampler.PriorDistributionSampler(self)
+
+  @functools.cached_property
+  def posterior_sampler_callable(
+      self,
+  ) -> posterior_sampler.PosteriorMCMCSampler:
+    """A `PosteriorMCMCSampler` callable bound to this model."""
+    return posterior_sampler.PosteriorMCMCSampler(self)
+
   def expand_selected_time_dims(
       self,
       start_date: tc.Date | None = None,
@@ -565,9 +537,10 @@ class Meridian:
     self._validate_injected_inference_data_group_coord(
         inference_data, group, constants.TIME, self.n_times
     )
-    self._validate_injected_inference_data_group_coord(
-        inference_data, group, constants.SIGMA_DIM, self._sigma_shape
-    )
+    if not self.model_spec.unique_sigma_for_each_geo:
+      self._validate_injected_inference_data_group_coord(
+          inference_data, group, constants.SIGMA_DIM, self._sigma_shape
+      )
     self._validate_injected_inference_data_group_coord(
         inference_data,
         group,
@@ -720,7 +693,7 @@ class Meridian:
       raise ValueError(
           f"Custom priors should be set on `{constants.MROI_M}` and"
           f" `{constants.MROI_RF}` when KPI is non-revenue and revenue per kpi"
-          f" data is missing."
+          " data is missing."
       )
 
   def _validate_geo_invariants(self):
@@ -955,143 +928,6 @@ class Meridian:
 
     return rf_out
 
-  def _get_roi_prior_beta_m_value(
-      self,
-      alpha_m: tf.Tensor,
-      beta_gm_dev: tf.Tensor,
-      ec_m: tf.Tensor,
-      eta_m: tf.Tensor,
-      roi_or_mroi_m: tf.Tensor,
-      slope_m: tf.Tensor,
-      media_transformed: tf.Tensor,
-  ) -> tf.Tensor:
-    """Returns a tensor to be used in `beta_m`."""
-    # The `roi_or_mroi_m` parameter represents either ROI or mROI. For reach &
-    # frequency channels, marginal ROI priors are defined as "mROI by reach",
-    # which is equivalent to ROI.
-    media_spend = self.media_tensors.media_spend
-    media_spend_counterfactual = self.media_tensors.media_spend_counterfactual
-    media_counterfactual_scaled = self.media_tensors.media_counterfactual_scaled
-    # If we got here, then we should already have media tensors derived from
-    # non-None InputData.media data.
-    assert media_spend is not None
-    assert media_spend_counterfactual is not None
-    assert media_counterfactual_scaled is not None
-
-    # Use absolute value here because this difference will be negative for
-    # marginal ROI priors.
-    inc_revenue_m = roi_or_mroi_m * tf.reduce_sum(
-        tf.abs(media_spend - media_spend_counterfactual),
-        range(media_spend.ndim - 1),
-    )
-
-    if (
-        self.model_spec.roi_calibration_period is None
-        and self.model_spec.paid_media_prior_type
-        == constants.PAID_MEDIA_PRIOR_TYPE_ROI
-    ):
-      # We can skip the adstock/hill computation step in this case.
-      media_counterfactual_transformed = tf.zeros_like(media_transformed)
-    else:
-      media_counterfactual_transformed = self.adstock_hill_media(
-          media=media_counterfactual_scaled,
-          alpha=alpha_m,
-          ec=ec_m,
-          slope=slope_m,
-      )
-
-    revenue_per_kpi = self.revenue_per_kpi
-    if self.input_data.revenue_per_kpi is None:
-      revenue_per_kpi = tf.ones([self.n_geos, self.n_times], dtype=tf.float32)
-    # Note: use absolute value here because this difference will be negative for
-    # marginal ROI priors.
-    media_contrib_gm = tf.einsum(
-        "...gtm,g,,gt->...gm",
-        tf.abs(media_transformed - media_counterfactual_transformed),
-        self.population,
-        self.kpi_transformer.population_scaled_stdev,
-        revenue_per_kpi,
-    )
-
-    if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL:
-      media_contrib_m = tf.einsum("...gm->...m", media_contrib_gm)
-      random_effect_m = tf.einsum(
-          "...m,...gm,...gm->...m", eta_m, beta_gm_dev, media_contrib_gm
-      )
-      return (inc_revenue_m - random_effect_m) / media_contrib_m
-    else:
-      # For log_normal, beta_m and eta_m are not mean & std.
-      # The parameterization is beta_gm ~ exp(beta_m + eta_m * N(0, 1)).
-      random_effect_m = tf.einsum(
-          "...gm,...gm->...m",
-          tf.math.exp(beta_gm_dev * eta_m[..., tf.newaxis, :]),
-          media_contrib_gm,
-      )
-    return tf.math.log(inc_revenue_m) - tf.math.log(random_effect_m)
-
-  def _get_roi_prior_beta_rf_value(
-      self,
-      alpha_rf: tf.Tensor,
-      beta_grf_dev: tf.Tensor,
-      ec_rf: tf.Tensor,
-      eta_rf: tf.Tensor,
-      roi_or_mroi_rf: tf.Tensor,
-      slope_rf: tf.Tensor,
-      rf_transformed: tf.Tensor,
-  ) -> tf.Tensor:
-    """Returns a tensor to be used in `beta_rf`."""
-    rf_spend = self.rf_tensors.rf_spend
-    rf_spend_counterfactual = self.rf_tensors.rf_spend_counterfactual
-    reach_counterfactual_scaled = self.rf_tensors.reach_counterfactual_scaled
-    frequency = self.rf_tensors.frequency
-    # If we got here, then we should already have RF media tensors derived from
-    # non-None InputData.reach data.
-    assert rf_spend is not None
-    assert rf_spend_counterfactual is not None
-    assert reach_counterfactual_scaled is not None
-    assert frequency is not None
-
-    inc_revenue_rf = roi_or_mroi_rf * tf.reduce_sum(
-        rf_spend - rf_spend_counterfactual,
-        range(rf_spend.ndim - 1),
-    )
-    if self.model_spec.rf_roi_calibration_period is not None:
-      rf_counterfactual_transformed = self.adstock_hill_rf(
-          reach=reach_counterfactual_scaled,
-          frequency=frequency,
-          alpha=alpha_rf,
-          ec=ec_rf,
-          slope=slope_rf,
-      )
-    else:
-      rf_counterfactual_transformed = tf.zeros_like(rf_transformed)
-    revenue_per_kpi = self.revenue_per_kpi
-    if self.input_data.revenue_per_kpi is None:
-      revenue_per_kpi = tf.ones([self.n_geos, self.n_times], dtype=tf.float32)
-
-    media_contrib_grf = tf.einsum(
-        "...gtm,g,,gt->...gm",
-        rf_transformed - rf_counterfactual_transformed,
-        self.population,
-        self.kpi_transformer.population_scaled_stdev,
-        revenue_per_kpi,
-    )
-    if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL:
-      media_contrib_rf = tf.einsum("...gm->...m", media_contrib_grf)
-      random_effect_rf = tf.einsum(
-          "...m,...gm,...gm->...m", eta_rf, beta_grf_dev, media_contrib_grf
-      )
-      return (inc_revenue_rf - random_effect_rf) / media_contrib_rf
-    else:
-      # For log_normal, beta_rf and eta_rf are not mean & std.
-      # The parameterization is beta_grf ~ exp(beta_rf + eta_rf * N(0, 1)).
-      random_effect_rf = tf.einsum(
-          "...gm,...gm->...m",
-          tf.math.exp(beta_grf_dev * eta_rf[..., tf.newaxis, :]),
-          media_contrib_grf,
-      )
-      return tf.math.log(inc_revenue_rf) - tf.math.log(random_effect_rf)
-
   def populate_cached_properties(self):
     """Eagerly activates all cached properties.
 
@@ -1111,301 +947,7 @@ class Meridian:
     for attr in cached_properties:
       _ = getattr(self, attr)
 
-  def _get_joint_dist_unpinned(self) -> tfp.distributions.Distribution:
-    """Returns JointDistributionCoroutineAutoBatched function for MCMC."""
-
-    self.populate_cached_properties()
-
-    # This lists all the derived properties and states of this Meridian object
-    # that are referenced by the joint distribution coroutine.
-    # That is, these are the list of captured parameters.
-    prior_broadcast = self.prior_broadcast
-    baseline_geo_idx = self.baseline_geo_idx
-    knot_info = self.knot_info
-    n_geos = self.n_geos
-    n_times = self.n_times
-    n_media_channels = self.n_media_channels
-    n_rf_channels = self.n_rf_channels
-    n_organic_media_channels = self.n_organic_media_channels
-    n_organic_rf_channels = self.n_organic_rf_channels
-    n_controls = self.n_controls
-    n_non_media_channels = self.n_non_media_channels
-    holdout_id = self.holdout_id
-    media_tensors = self.media_tensors
-    rf_tensors = self.rf_tensors
-    organic_media_tensors = self.organic_media_tensors
-    organic_rf_tensors = self.organic_rf_tensors
-    controls_scaled = self.controls_scaled
-    non_media_treatments_scaled = self.non_media_treatments_scaled
-    media_effects_dist = self.media_effects_dist
-    adstock_hill_media_fn = self.adstock_hill_media
-    adstock_hill_rf_fn = self.adstock_hill_rf
-    get_roi_prior_beta_m_value_fn = self._get_roi_prior_beta_m_value
-    get_roi_prior_beta_rf_value_fn = self._get_roi_prior_beta_rf_value
-
-    # TODO: Extract this coroutine to be unittestable on its own.
-    # This MCMC sampling technique is complex enough to have its own abstraction
-    # and testable API, rather than being embedded as a private method in the
-    # Meridian class.
-    @tfp.distributions.JointDistributionCoroutineAutoBatched
-    def joint_dist_unpinned():
-      # Sample directly from prior.
-      knot_values = yield prior_broadcast.knot_values
-      gamma_c = yield prior_broadcast.gamma_c
-      xi_c = yield prior_broadcast.xi_c
-      sigma = yield prior_broadcast.sigma
-
-      tau_g_excl_baseline = yield tfp.distributions.Sample(
-          prior_broadcast.tau_g_excl_baseline,
-          name=constants.TAU_G_EXCL_BASELINE,
-      )
-      tau_g = yield _get_tau_g(
-          tau_g_excl_baseline=tau_g_excl_baseline,
-          baseline_geo_idx=baseline_geo_idx,
-      )
-      mu_t = yield tfp.distributions.Deterministic(
-          tf.einsum(
-              "k,kt->t",
-              knot_values,
-              tf.convert_to_tensor(knot_info.weights),
-          ),
-          name=constants.MU_T,
-      )
-
-      tau_gt = tau_g[:, tf.newaxis] + mu_t
-      combined_media_transformed = tf.zeros(
-          shape=(n_geos, n_times, 0), dtype=tf.float32
-      )
-      combined_beta = tf.zeros(shape=(n_geos, 0), dtype=tf.float32)
-      if media_tensors.media is not None:
-        alpha_m = yield prior_broadcast.alpha_m
-        ec_m = yield prior_broadcast.ec_m
-        eta_m = yield prior_broadcast.eta_m
-        slope_m = yield prior_broadcast.slope_m
-        beta_gm_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
-            [n_geos, n_media_channels],
-            name=constants.BETA_GM_DEV,
-        )
-        media_transformed = adstock_hill_media_fn(
-            media=media_tensors.media_scaled,
-            alpha=alpha_m,
-            ec=ec_m,
-            slope=slope_m,
-        )
-        prior_type = self.model_spec.paid_media_prior_type
-        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
-          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-            roi_or_mroi_m = yield prior_broadcast.roi_m
-          else:
-            roi_or_mroi_m = yield prior_broadcast.mroi_m
-          beta_m_value = get_roi_prior_beta_m_value_fn(
-              alpha_m,
-              beta_gm_dev,
-              ec_m,
-              eta_m,
-              roi_or_mroi_m,
-              slope_m,
-              media_transformed,
-          )
-          beta_m = yield tfp.distributions.Deterministic(
-              beta_m_value, name=constants.BETA_M
-          )
-        else:
-          beta_m = yield prior_broadcast.beta_m
-
-        beta_eta_combined = beta_m + eta_m * beta_gm_dev
-        beta_gm_value = (
-            beta_eta_combined
-            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
-        )
-        beta_gm = yield tfp.distributions.Deterministic(
-            beta_gm_value, name=constants.BETA_GM
-        )
-        combined_media_transformed = tf.concat(
-            [combined_media_transformed, media_transformed], axis=-1
-        )
-        combined_beta = tf.concat([combined_beta, beta_gm], axis=-1)
-
-      if rf_tensors.reach is not None:
-        alpha_rf = yield prior_broadcast.alpha_rf
-        ec_rf = yield prior_broadcast.ec_rf
-        eta_rf = yield prior_broadcast.eta_rf
-        slope_rf = yield prior_broadcast.slope_rf
-        beta_grf_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
-            [n_geos, n_rf_channels],
-            name=constants.BETA_GRF_DEV,
-        )
-        rf_transformed = adstock_hill_rf_fn(
-            reach=rf_tensors.reach_scaled,
-            frequency=rf_tensors.frequency,
-            alpha=alpha_rf,
-            ec=ec_rf,
-            slope=slope_rf,
-        )
-
-        prior_type = self.model_spec.paid_media_prior_type
-        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
-          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-            roi_or_mroi_rf = yield prior_broadcast.roi_rf
-          else:
-            roi_or_mroi_rf = yield prior_broadcast.mroi_rf
-          beta_rf_value = get_roi_prior_beta_rf_value_fn(
-              alpha_rf,
-              beta_grf_dev,
-              ec_rf,
-              eta_rf,
-              roi_or_mroi_rf,
-              slope_rf,
-              rf_transformed,
-          )
-          beta_rf = yield tfp.distributions.Deterministic(
-              beta_rf_value,
-              name=constants.BETA_RF,
-          )
-        else:
-          beta_rf = yield prior_broadcast.beta_rf
-
-        beta_eta_combined = beta_rf + eta_rf * beta_grf_dev
-        beta_grf_value = (
-            beta_eta_combined
-            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
-        )
-        beta_grf = yield tfp.distributions.Deterministic(
-            beta_grf_value, name=constants.BETA_GRF
-        )
-        combined_media_transformed = tf.concat(
-            [combined_media_transformed, rf_transformed], axis=-1
-        )
-        combined_beta = tf.concat([combined_beta, beta_grf], axis=-1)
-
-      if organic_media_tensors.organic_media is not None:
-        alpha_om = yield prior_broadcast.alpha_om
-        ec_om = yield prior_broadcast.ec_om
-        eta_om = yield prior_broadcast.eta_om
-        slope_om = yield prior_broadcast.slope_om
-        beta_gom_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
-            [n_geos, n_organic_media_channels],
-            name=constants.BETA_GOM_DEV,
-        )
-        organic_media_transformed = adstock_hill_media_fn(
-            media=organic_media_tensors.organic_media_scaled,
-            alpha=alpha_om,
-            ec=ec_om,
-            slope=slope_om,
-        )
-        beta_om = yield prior_broadcast.beta_om
-
-        beta_eta_combined = beta_om + eta_om * beta_gom_dev
-        beta_gom_value = (
-            beta_eta_combined
-            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
-        )
-        beta_gom = yield tfp.distributions.Deterministic(
-            beta_gom_value, name=constants.BETA_GOM
-        )
-        combined_media_transformed = tf.concat(
-            [combined_media_transformed, organic_media_transformed], axis=-1
-        )
-        combined_beta = tf.concat([combined_beta, beta_gom], axis=-1)
-
-      if organic_rf_tensors.organic_reach is not None:
-        alpha_orf = yield prior_broadcast.alpha_orf
-        ec_orf = yield prior_broadcast.ec_orf
-        eta_orf = yield prior_broadcast.eta_orf
-        slope_orf = yield prior_broadcast.slope_orf
-        beta_gorf_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
-            [n_geos, n_organic_rf_channels],
-            name=constants.BETA_GORF_DEV,
-        )
-        organic_rf_transformed = adstock_hill_rf_fn(
-            reach=organic_rf_tensors.organic_reach_scaled,
-            frequency=organic_rf_tensors.organic_frequency,
-            alpha=alpha_orf,
-            ec=ec_orf,
-            slope=slope_orf,
-        )
-
-        beta_orf = yield prior_broadcast.beta_orf
-
-        beta_eta_combined = beta_orf + eta_orf * beta_gorf_dev
-        beta_gorf_value = (
-            beta_eta_combined
-            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
-        )
-        beta_gorf = yield tfp.distributions.Deterministic(
-            beta_gorf_value, name=constants.BETA_GORF
-        )
-        combined_media_transformed = tf.concat(
-            [combined_media_transformed, organic_rf_transformed], axis=-1
-        )
-        combined_beta = tf.concat([combined_beta, beta_gorf], axis=-1)
-
-      sigma_gt = tf.transpose(tf.broadcast_to(sigma, [n_times, n_geos]))
-      gamma_gc_dev = yield tfp.distributions.Sample(
-          tfp.distributions.Normal(0, 1),
-          [n_geos, n_controls],
-          name=constants.GAMMA_GC_DEV,
-      )
-      gamma_gc = yield tfp.distributions.Deterministic(
-          gamma_c + xi_c * gamma_gc_dev, name=constants.GAMMA_GC
-      )
-      y_pred_combined_media = (
-          tau_gt
-          + tf.einsum("gtm,gm->gt", combined_media_transformed, combined_beta)
-          + tf.einsum("gtc,gc->gt", controls_scaled, gamma_gc)
-      )
-
-      if self.non_media_treatments is not None:
-        gamma_n = yield prior_broadcast.gamma_n
-        xi_n = yield prior_broadcast.xi_n
-        gamma_gn_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
-            [n_geos, n_non_media_channels],
-            name=constants.GAMMA_GN_DEV,
-        )
-        gamma_gn = yield tfp.distributions.Deterministic(
-            gamma_n + xi_n * gamma_gn_dev, name=constants.GAMMA_GN
-        )
-        y_pred = y_pred_combined_media + tf.einsum(
-            "gtn,gn->gt", non_media_treatments_scaled, gamma_gn
-        )
-      else:
-        y_pred = y_pred_combined_media
-
-      # If there are any holdout observations, the holdout KPI values will
-      # be replaced with zeros using `experimental_pin`. For these
-      # observations, we set the posterior mean equal to zero and standard
-      # deviation to `1/sqrt(2pi)`, so the log-density is 0 regardless of the
-      # sampled posterior parameter values.
-      if holdout_id is not None:
-        y_pred_holdout = tf.where(holdout_id, 0.0, y_pred)
-        test_sd = tf.cast(1.0 / np.sqrt(2.0 * np.pi), tf.float32)
-        sigma_gt_holdout = tf.where(holdout_id, test_sd, sigma_gt)
-        yield tfp.distributions.Normal(
-            y_pred_holdout, sigma_gt_holdout, name="y"
-        )
-      else:
-        yield tfp.distributions.Normal(y_pred, sigma_gt, name="y")
-
-    return joint_dist_unpinned
-
-  def _get_joint_dist(self) -> tfp.distributions.Distribution:
-    y = (
-        tf.where(self.holdout_id, 0.0, self.kpi_scaled)
-        if self.holdout_id is not None
-        else self.kpi_scaled
-    )
-    return self._get_joint_dist_unpinned().experimental_pin(y=y)
-
-  def _create_inference_data_coords(
+  def create_inference_data_coords(
       self, n_chains: int, n_draws: int
   ) -> Mapping[str, np.ndarray | Sequence[str]]:
     """Creates data coordinates for inference data."""
@@ -1449,7 +991,7 @@ class Meridian:
         constants.ORGANIC_RF_CHANNEL: organic_rf_channel_values,
     }
 
-  def _create_inference_data_dims(self) -> Mapping[str, Sequence[str]]:
+  def create_inference_data_dims(self) -> Mapping[str, Sequence[str]]:
     inference_dims = dict(constants.INFERENCE_DIMS)
     if self.unique_sigma_for_each_geo:
       inference_dims[constants.SIGMA] = [constants.GEO]
@@ -1461,412 +1003,18 @@ class Meridian:
         for param, dims in inference_dims.items()
     }
 
-  def _sample_media_priors(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.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).
-
-    Returns:
-      A mapping of media parameter names to a tensor of shape [n_draws, n_geos,
-      n_media_channels] or [n_draws, n_media_channels] containing the
-      samples.
-    """
-    prior = self.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),
-    }
-    beta_gm_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_media_channels],
-        name=constants.BETA_GM_DEV,
-    ).sample(**sample_kwargs)
-    media_transformed = self.adstock_hill_media(
-        media=self.media_tensors.media_scaled,
-        alpha=media_vars[constants.ALPHA_M],
-        ec=media_vars[constants.EC_M],
-        slope=media_vars[constants.SLOPE_M],
-    )
-
-    prior_type = self.model_spec.paid_media_prior_type
-    if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-      roi_m = prior.roi_m.sample(**sample_kwargs)
-      beta_m_value = self._get_roi_prior_beta_m_value(
-          beta_gm_dev=beta_gm_dev,
-          media_transformed=media_transformed,
-          roi_or_mroi_m=roi_m,
-          **media_vars,
-      )
-      media_vars[constants.ROI_M] = roi_m
-      media_vars[constants.BETA_M] = tfp.distributions.Deterministic(
-          beta_m_value, name=constants.BETA_M
-      ).sample()
-    elif prior_type == constants.PAID_MEDIA_PRIOR_TYPE_MROI:
-      mroi_m = prior.mroi_m.sample(**sample_kwargs)
-      beta_m_value = self._get_roi_prior_beta_m_value(
-          beta_gm_dev=beta_gm_dev,
-          media_transformed=media_transformed,
-          roi_or_mroi_m=mroi_m,
-          **media_vars,
-      )
-      media_vars[constants.MROI_M] = mroi_m
-      media_vars[constants.BETA_M] = tfp.distributions.Deterministic(
-          beta_m_value, name=constants.BETA_M
-      ).sample()
-    else:
-      media_vars[constants.BETA_M] = prior.beta_m.sample(**sample_kwargs)
-
-    beta_eta_combined = (
-        media_vars[constants.BETA_M][..., tf.newaxis, :]
-        + media_vars[constants.ETA_M][..., tf.newaxis, :] * beta_gm_dev
-    )
-    beta_gm_value = (
-        beta_eta_combined
-        if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
-    )
-    media_vars[constants.BETA_GM] = tfp.distributions.Deterministic(
-        beta_gm_value, name=constants.BETA_GM
-    ).sample()
-
-    return media_vars
-
-  def _sample_rf_priors(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.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).
-
-    Returns:
-      A mapping of RF parameter names to a tensor of shape [n_draws, n_geos,
-      n_rf_channels] or [n_draws, n_rf_channels] containing the samples.
-    """
-    prior = self.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),
-    }
-    beta_grf_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_rf_channels],
-        name=constants.BETA_GRF_DEV,
-    ).sample(**sample_kwargs)
-    rf_transformed = self.adstock_hill_rf(
-        reach=self.rf_tensors.reach_scaled,
-        frequency=self.rf_tensors.frequency,
-        alpha=rf_vars[constants.ALPHA_RF],
-        ec=rf_vars[constants.EC_RF],
-        slope=rf_vars[constants.SLOPE_RF],
-    )
-
-    prior_type = self.model_spec.paid_media_prior_type
-    if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-      roi_rf = prior.roi_rf.sample(**sample_kwargs)
-      beta_rf_value = self._get_roi_prior_beta_rf_value(
-          beta_grf_dev=beta_grf_dev,
-          rf_transformed=rf_transformed,
-          roi_or_mroi_rf=roi_rf,
-          **rf_vars,
-      )
-      rf_vars[constants.ROI_RF] = roi_rf
-      rf_vars[constants.BETA_RF] = tfp.distributions.Deterministic(
-          beta_rf_value,
-          name=constants.BETA_RF,
-      ).sample()
-    elif prior_type == constants.PAID_MEDIA_PRIOR_TYPE_MROI:
-      mroi_rf = prior.mroi_rf.sample(**sample_kwargs)
-      beta_rf_value = self._get_roi_prior_beta_rf_value(
-          beta_grf_dev=beta_grf_dev,
-          rf_transformed=rf_transformed,
-          roi_or_mroi_rf=mroi_rf,
-          **rf_vars,
-      )
-      rf_vars[constants.MROI_RF] = mroi_rf
-      rf_vars[constants.BETA_RF] = tfp.distributions.Deterministic(
-          beta_rf_value,
-          name=constants.BETA_RF,
-      ).sample()
-    else:
-      rf_vars[constants.BETA_RF] = prior.beta_rf.sample(**sample_kwargs)
-
-    beta_eta_combined = (
-        rf_vars[constants.BETA_RF][..., tf.newaxis, :]
-        + rf_vars[constants.ETA_RF][..., tf.newaxis, :] * beta_grf_dev
-    )
-    beta_grf_value = (
-        beta_eta_combined
-        if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
-    )
-    rf_vars[constants.BETA_GRF] = tfp.distributions.Deterministic(
-        beta_grf_value, name=constants.BETA_GRF
-    ).sample()
-
-    return rf_vars
-
-  def _sample_organic_media_priors(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.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).
-
-    Returns:
-      A mapping of organic media parameter names to a tensor of shape [n_draws,
-      n_geos, n_organic_media_channels] or [n_draws, n_organic_media_channels]
-      containing the samples.
-    """
-    prior = self.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),
-    }
-    beta_gom_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_organic_media_channels],
-        name=constants.BETA_GOM_DEV,
-    ).sample(**sample_kwargs)
-
-    organic_media_vars[constants.BETA_OM] = prior.beta_om.sample(
-        **sample_kwargs
-    )
-
-    beta_eta_combined = (
-        organic_media_vars[constants.BETA_OM][..., tf.newaxis, :]
-        + organic_media_vars[constants.ETA_OM][..., tf.newaxis, :]
-        * beta_gom_dev
-    )
-    beta_gom_value = (
-        beta_eta_combined
-        if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
-    )
-    organic_media_vars[constants.BETA_GOM] = tfp.distributions.Deterministic(
-        beta_gom_value, name=constants.BETA_GOM
-    ).sample()
-
-    return organic_media_vars
-
-  def _sample_organic_rf_priors(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.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).
-
-    Returns:
-      A mapping of organic RF parameter names to a tensor of shape [n_draws,
-      n_geos, n_organic_rf_channels] or [n_draws, n_organic_rf_channels]
-      containing the samples.
-    """
-    prior = self.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),
-    }
-    beta_gorf_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_organic_rf_channels],
-        name=constants.BETA_GORF_DEV,
-    ).sample(**sample_kwargs)
-
-    organic_rf_vars[constants.BETA_ORF] = prior.beta_orf.sample(**sample_kwargs)
-
-    beta_eta_combined = (
-        organic_rf_vars[constants.BETA_ORF][..., tf.newaxis, :]
-        + organic_rf_vars[constants.ETA_ORF][..., tf.newaxis, :] * beta_gorf_dev
-    )
-    beta_gorf_value = (
-        beta_eta_combined
-        if self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
-    )
-    organic_rf_vars[constants.BETA_GORF] = tfp.distributions.Deterministic(
-        beta_gorf_value, name=constants.BETA_GORF
-    ).sample()
-
-    return organic_rf_vars
-
-  def _sample_non_media_treatments_priors(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.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).
-
-    Returns:
-      A mapping of non-media treatment parameter names to a tensor of shape
-      [n_draws,
-      n_geos, n_non_media_channels] or [n_draws, n_non_media_channels]
-      containing the samples.
-    """
-    prior = self.prior_broadcast
-    sample_shape = [1, n_draws]
-    sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
-    non_media_treatments_vars = {
-        constants.GAMMA_N: prior.gamma_n.sample(**sample_kwargs),
-        constants.XI_N: prior.xi_n.sample(**sample_kwargs),
-    }
-    gamma_gn_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_non_media_channels],
-        name=constants.GAMMA_GN_DEV,
-    ).sample(**sample_kwargs)
-    non_media_treatments_vars[constants.GAMMA_GN] = (
-        tfp.distributions.Deterministic(
-            non_media_treatments_vars[constants.GAMMA_N][..., tf.newaxis, :]
-            + non_media_treatments_vars[constants.XI_N][..., tf.newaxis, :]
-            * gamma_gn_dev,
-            name=constants.GAMMA_GN,
-        ).sample()
-    )
-    return non_media_treatments_vars
-
-  def _sample_prior_fn(
-      self,
-      n_draws: int,
-      seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
-    """Returns a mapping of prior parameters to tensors of the samples."""
-    # For stateful sampling, the random seed must be set to ensure that any
-    # random numbers that are generated are deterministic.
-    if seed is not None:
-      tf.keras.utils.set_random_seed(1)
-    prior = self.prior_broadcast
-    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)
-    base_vars = {
-        constants.KNOT_VALUES: prior.knot_values.sample(**sample_kwargs),
-        constants.GAMMA_C: prior.gamma_c.sample(**sample_kwargs),
-        constants.XI_C: prior.xi_c.sample(**sample_kwargs),
-        constants.SIGMA: prior.sigma.sample(**sample_kwargs),
-        constants.TAU_G: _get_tau_g(
-            tau_g_excl_baseline=tau_g_excl_baseline,
-            baseline_geo_idx=self.baseline_geo_idx,
-        ).sample(),
-    }
-    base_vars[constants.MU_T] = tfp.distributions.Deterministic(
-        tf.einsum(
-            "...k,kt->...t",
-            base_vars[constants.KNOT_VALUES],
-            tf.convert_to_tensor(self.knot_info.weights),
-        ),
-        name=constants.MU_T,
-    ).sample()
-
-    gamma_gc_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
-        [self.n_geos, self.n_controls],
-        name=constants.GAMMA_GC_DEV,
-    ).sample(**sample_kwargs)
-    base_vars[constants.GAMMA_GC] = tfp.distributions.Deterministic(
-        base_vars[constants.GAMMA_C][..., tf.newaxis, :]
-        + base_vars[constants.XI_C][..., tf.newaxis, :] * gamma_gc_dev,
-        name=constants.GAMMA_GC,
-    ).sample()
-
-    media_vars = (
-        self._sample_media_priors(n_draws, seed)
-        if self.media_tensors.media is not None
-        else {}
-    )
-    rf_vars = (
-        self._sample_rf_priors(n_draws, seed)
-        if self.rf_tensors.reach is not None
-        else {}
-    )
-    organic_media_vars = (
-        self._sample_organic_media_priors(n_draws, seed)
-        if self.organic_media_tensors.organic_media is not None
-        else {}
-    )
-    organic_rf_vars = (
-        self._sample_organic_rf_priors(n_draws, seed)
-        if self.organic_rf_tensors.organic_reach is not None
-        else {}
-    )
-    non_media_treatments_vars = (
-        self._sample_non_media_treatments_priors(n_draws, seed)
-        if self.non_media_treatments_scaled is not None
-        else {}
-    )
-
-    return (
-        base_vars
-        | media_vars
-        | rf_vars
-        | organic_media_vars
-        | organic_rf_vars
-        | non_media_treatments_vars
-    )
-
   def sample_prior(self, n_draws: int, seed: int | None = None):
     """Draws samples from the prior distributions.
 
+    Drawn samples are merged into this model's Arviz `inference_data` property.
+
     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).
     """
-    prior_draws = self._sample_prior_fn(n_draws, seed=seed)
-    # Create Arviz InferenceData for prior draws.
-    prior_coords = self._create_inference_data_coords(1, n_draws)
-    prior_dims = self._create_inference_data_dims()
-    prior_inference_data = az.convert_to_inference_data(
-        prior_draws, coords=prior_coords, dims=prior_dims, group=constants.PRIOR
-    )
+    prior_inference_data = self.prior_sampler_callable(n_draws, seed)
     self.inference_data.extend(prior_inference_data, join="right")
 
   def sample_posterior(
@@ -1890,6 +1038,8 @@ class Meridian:
     For more information about the arguments, see [`windowed_adaptive_nuts`]
     (https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/mcmc/windowed_adaptive_nuts).
 
+    Drawn samples are merged into this model's Arviz `inference_data` property.
+
     Args:
       n_chains: Number of MCMC chains. Given a sequence of integers,
         `windowed_adaptive_nuts` will be called once for each element. The
@@ -1943,112 +1093,20 @@ class Meridian:
         [ResourceExhaustedError when running Meridian.sample_posterior]
         (https://developers.google.com/meridian/docs/advanced-modeling/model-debugging#gpu-oom-error).
     """
-    seed = tfp.random.sanitize_seed(seed) if seed else None
-    n_chains_list = [n_chains] if isinstance(n_chains, int) else n_chains
-    total_chains = np.sum(n_chains_list)
-
-    states = []
-    traces = []
-    for n_chains_batch in n_chains_list:
-      try:
-        mcmc = _xla_windowed_adaptive_nuts(
-            n_draws=n_burnin + n_keep,
-            joint_dist=self._get_joint_dist(),
-            n_chains=n_chains_batch,
-            num_adaptation_steps=n_adapt,
-            current_state=current_state,
-            init_step_size=init_step_size,
-            dual_averaging_kwargs=dual_averaging_kwargs,
-            max_tree_depth=max_tree_depth,
-            max_energy_diff=max_energy_diff,
-            unrolled_leapfrog_steps=unrolled_leapfrog_steps,
-            parallel_iterations=parallel_iterations,
-            seed=seed,
-            **pins,
-        )
-      except tf.errors.ResourceExhaustedError as error:
-        raise MCMCOOMError(
-            "ERROR: Out of memory. Try reducing `n_keep` or pass a list of"
-            " integers as `n_chains` to sample chains serially (see"
-            " https://developers.google.com/meridian/docs/advanced-modeling/model-debugging#gpu-oom-error)"
-        ) from error
-      states.append(mcmc.all_states._asdict())
-      traces.append(mcmc.trace)
-
-    mcmc_states = {
-        k: tf.einsum(
-            "ij...->ji...",
-            tf.concat([state[k] for state in states], axis=1)[n_burnin:, ...],
-        )
-        for k in states[0].keys()
-        if k not in constants.UNSAVED_PARAMETERS
-    }
-    # Create Arviz InferenceData for posterior draws.
-    posterior_coords = self._create_inference_data_coords(total_chains, n_keep)
-    posterior_dims = self._create_inference_data_dims()
-    infdata_posterior = az.convert_to_inference_data(
-        mcmc_states, coords=posterior_coords, dims=posterior_dims
-    )
-
-    # Save trace metrics in InferenceData.
-    mcmc_trace = {}
-    for k in traces[0].keys():
-      if k not in constants.IGNORED_TRACE_METRICS:
-        mcmc_trace[k] = tf.concat(
-            [
-                tf.broadcast_to(
-                    tf.transpose(trace[k][n_burnin:, ...]),
-                    [n_chains_list[i], n_keep],
-                )
-                for i, trace in enumerate(traces)
-            ],
-            axis=0,
-        )
-
-    trace_coords = {
-        constants.CHAIN: np.arange(total_chains),
-        constants.DRAW: np.arange(n_keep),
-    }
-    trace_dims = {
-        k: [constants.CHAIN, constants.DRAW] for k in mcmc_trace.keys()
-    }
-    infdata_trace = az.convert_to_inference_data(
-        mcmc_trace, coords=trace_coords, dims=trace_dims, group="trace"
-    )
-
-    # Create Arviz InferenceData for divergent transitions and other sampling
-    # statistics. Note that InferenceData has a different naming convention
-    # than Tensorflow, and only certain variables are recongnized.
-    # https://arviz-devs.github.io/arviz/schema/schema.html#sample-stats
-    # The list of values returned by windowed_adaptive_nuts() is the following:
-    # 'step_size', 'tune', 'target_log_prob', 'diverging', 'accept_ratio',
-    # 'variance_scaling', 'n_steps', 'is_accepted'.
-
-    sample_stats = {
-        constants.SAMPLE_STATS_METRICS[k]: v
-        for k, v in mcmc_trace.items()
-        if k in constants.SAMPLE_STATS_METRICS
-    }
-    sample_stats_dims = {
-        constants.SAMPLE_STATS_METRICS[k]: v
-        for k, v in trace_dims.items()
-        if k in constants.SAMPLE_STATS_METRICS
-    }
-    # Tensorflow does not include a "draw" dimension on step size metric if same
-    # step size is used for all chains. Step size must be broadcast to the
-    # correct shape.
-    sample_stats[constants.STEP_SIZE] = tf.broadcast_to(
-        sample_stats[constants.STEP_SIZE], [total_chains, n_keep]
-    )
-    sample_stats_dims[constants.STEP_SIZE] = [constants.CHAIN, constants.DRAW]
-    infdata_sample_stats = az.convert_to_inference_data(
-        sample_stats,
-        coords=trace_coords,
-        dims=sample_stats_dims,
-        group="sample_stats",
-    )
-    posterior_inference_data = az.concat(
-        infdata_posterior, infdata_trace, infdata_sample_stats
+    posterior_inference_data = self.posterior_sampler_callable(
+        n_chains,
+        n_adapt,
+        n_burnin,
+        n_keep,
+        current_state,
+        init_step_size,
+        dual_averaging_kwargs,
+        max_tree_depth,
+        max_energy_diff,
+        unrolled_leapfrog_steps,
+        parallel_iterations,
+        seed,
+        **pins,
     )
     self.inference_data.extend(posterior_inference_data, join="right")