google-meridian 1.1.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

meridian/model/posterior_sampler.py

@@ -18,10 +18,9 @@ from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING
 
 import arviz as az
+from meridian import backend
 from meridian import constants
 import numpy as np
-import tensorflow as tf
-import tensorflow_probability as tfp
 
 if TYPE_CHECKING:
   from meridian.model import model  # pylint: disable=g-bad-import-order,g-import-not-at-top
@@ -43,8 +42,8 @@ class MCMCOOMError(Exception):
 
 
 def _get_tau_g(
-    tau_g_excl_baseline: tf.Tensor, baseline_geo_idx: int
-) -> tfp.distributions.Distribution:
+    tau_g_excl_baseline: backend.Tensor, baseline_geo_idx: int
+) -> backend.tfd.Distribution:
   """Computes `tau_g` from `tau_g_excl_baseline`.
 
   This function computes `tau_g` by inserting a column of zeros at the
@@ -62,21 +61,21 @@ def _get_tau_g(
   """
   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 = backend.concatenate(
       [
           tau_g_excl_baseline[..., :baseline_geo_idx],
-          tf.zeros(shape, dtype=tau_g_excl_baseline.dtype),
+          backend.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")
+  return backend.tfd.Deterministic(tau_g, name="tau_g")
 
 
-@tf.function(autograph=False, jit_compile=True)
+@backend.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)
+  return backend.experimental.mcmc.windowed_adaptive_nuts(**kwargs)
 
 
 class PosteriorMCMCSampler:
@@ -89,7 +88,7 @@ class PosteriorMCMCSampler:
   def model(self) -> "model.Meridian":
     return self._meridian
 
-  def _get_joint_dist_unpinned(self) -> tfp.distributions.Distribution:
+  def _get_joint_dist_unpinned(self) -> backend.tfd.Distribution:
     """Returns a `JointDistributionCoroutineAutoBatched` function for MCMC."""
     mmm = self.model
     mmm.populate_cached_properties()
@@ -120,13 +119,13 @@ class PosteriorMCMCSampler:
     adstock_hill_rf_fn = mmm.adstock_hill_rf
     total_outcome = mmm.total_outcome
 
-    @tfp.distributions.JointDistributionCoroutineAutoBatched
+    @backend.tfd.JointDistributionCoroutineAutoBatched
     def joint_dist_unpinned():
       # Sample directly from prior.
       knot_values = yield prior_broadcast.knot_values
       sigma = yield prior_broadcast.sigma
 
-      tau_g_excl_baseline = yield tfp.distributions.Sample(
+      tau_g_excl_baseline = yield backend.tfd.Sample(
           prior_broadcast.tau_g_excl_baseline,
           name=constants.TAU_G_EXCL_BASELINE,
       )
@@ -134,27 +133,27 @@ class PosteriorMCMCSampler:
           tau_g_excl_baseline=tau_g_excl_baseline,
           baseline_geo_idx=baseline_geo_idx,
       )
-      mu_t = yield tfp.distributions.Deterministic(
-          tf.einsum(
+      mu_t = yield backend.tfd.Deterministic(
+          backend.einsum(
               "k,kt->t",
               knot_values,
-              tf.convert_to_tensor(knot_info.weights),
+              backend.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
+      tau_gt = tau_g[:, backend.newaxis] + mu_t
+      combined_media_transformed = backend.zeros(
+          shape=(n_geos, n_times, 0), dtype=backend.float32
       )
-      combined_beta = tf.zeros(shape=(n_geos, 0), dtype=tf.float32)
+      combined_beta = backend.zeros(shape=(n_geos, 0), dtype=backend.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),
+        beta_gm_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_media_channels],
             name=constants.BETA_GM_DEV,
         )
@@ -163,6 +162,7 @@ class PosteriorMCMCSampler:
             alpha=alpha_m,
             ec=ec_m,
             slope=slope_m,
+            decay_functions=mmm.adstock_decay_spec.media,
         )
         prior_type = mmm.model_spec.effective_media_prior_type
         if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
@@ -194,7 +194,7 @@ class PosteriorMCMCSampler:
               eta_x=eta_m,
               beta_gx_dev=beta_gm_dev,
           )
-          beta_m = yield tfp.distributions.Deterministic(
+          beta_m = yield backend.tfd.Deterministic(
               beta_m_value, name=constants.BETA_M
           )
 
@@ -202,23 +202,23 @@ class PosteriorMCMCSampler:
         beta_gm_value = (
             beta_eta_combined
             if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
+            else backend.exp(beta_eta_combined)
         )
-        beta_gm = yield tfp.distributions.Deterministic(
+        beta_gm = yield backend.tfd.Deterministic(
             beta_gm_value, name=constants.BETA_GM
         )
-        combined_media_transformed = tf.concat(
+        combined_media_transformed = backend.concatenate(
             [combined_media_transformed, media_transformed], axis=-1
         )
-        combined_beta = tf.concat([combined_beta, beta_gm], axis=-1)
+        combined_beta = backend.concatenate([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),
+        beta_grf_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_rf_channels],
             name=constants.BETA_GRF_DEV,
         )
@@ -228,6 +228,7 @@ class PosteriorMCMCSampler:
             alpha=alpha_rf,
             ec=ec_rf,
             slope=slope_rf,
+            decay_functions=mmm.adstock_decay_spec.rf,
         )
 
         prior_type = mmm.model_spec.effective_rf_prior_type
@@ -260,7 +261,7 @@ class PosteriorMCMCSampler:
               eta_x=eta_rf,
               beta_gx_dev=beta_grf_dev,
           )
-          beta_rf = yield tfp.distributions.Deterministic(
+          beta_rf = yield backend.tfd.Deterministic(
               beta_rf_value, name=constants.BETA_RF
           )
 
@@ -268,23 +269,23 @@ class PosteriorMCMCSampler:
         beta_grf_value = (
             beta_eta_combined
             if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
+            else backend.exp(beta_eta_combined)
         )
-        beta_grf = yield tfp.distributions.Deterministic(
+        beta_grf = yield backend.tfd.Deterministic(
             beta_grf_value, name=constants.BETA_GRF
         )
-        combined_media_transformed = tf.concat(
+        combined_media_transformed = backend.concatenate(
             [combined_media_transformed, rf_transformed], axis=-1
         )
-        combined_beta = tf.concat([combined_beta, beta_grf], axis=-1)
+        combined_beta = backend.concatenate([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),
+        beta_gom_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_organic_media_channels],
             name=constants.BETA_GOM_DEV,
         )
@@ -293,6 +294,7 @@ class PosteriorMCMCSampler:
             alpha=alpha_om,
             ec=ec_om,
             slope=slope_om,
+            decay_functions=mmm.adstock_decay_spec.organic_media,
         )
         prior_type = mmm.model_spec.organic_media_prior_type
         if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
@@ -307,7 +309,7 @@ class PosteriorMCMCSampler:
               eta_x=eta_om,
               beta_gx_dev=beta_gom_dev,
           )
-          beta_om = yield tfp.distributions.Deterministic(
+          beta_om = yield backend.tfd.Deterministic(
               beta_om_value, name=constants.BETA_OM
           )
         else:
@@ -317,23 +319,23 @@ class PosteriorMCMCSampler:
         beta_gom_value = (
             beta_eta_combined
             if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
+            else backend.exp(beta_eta_combined)
         )
-        beta_gom = yield tfp.distributions.Deterministic(
+        beta_gom = yield backend.tfd.Deterministic(
             beta_gom_value, name=constants.BETA_GOM
         )
-        combined_media_transformed = tf.concat(
+        combined_media_transformed = backend.concatenate(
             [combined_media_transformed, organic_media_transformed], axis=-1
         )
-        combined_beta = tf.concat([combined_beta, beta_gom], axis=-1)
+        combined_beta = backend.concatenate([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),
+        beta_gorf_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_organic_rf_channels],
             name=constants.BETA_GORF_DEV,
         )
@@ -343,6 +345,7 @@ class PosteriorMCMCSampler:
             alpha=alpha_orf,
             ec=ec_orf,
             slope=slope_orf,
+            decay_functions=mmm.adstock_decay_spec.organic_rf,
         )
 
         prior_type = mmm.model_spec.organic_rf_prior_type
@@ -358,7 +361,7 @@ class PosteriorMCMCSampler:
               eta_x=eta_orf,
               beta_gx_dev=beta_gorf_dev,
           )
-          beta_orf = yield tfp.distributions.Deterministic(
+          beta_orf = yield backend.tfd.Deterministic(
               beta_orf_value, name=constants.BETA_ORF
           )
         else:
@@ -368,18 +371,20 @@ class PosteriorMCMCSampler:
         beta_gorf_value = (
             beta_eta_combined
             if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-            else tf.math.exp(beta_eta_combined)
+            else backend.exp(beta_eta_combined)
         )
-        beta_gorf = yield tfp.distributions.Deterministic(
+        beta_gorf = yield backend.tfd.Deterministic(
             beta_gorf_value, name=constants.BETA_GORF
         )
-        combined_media_transformed = tf.concat(
+        combined_media_transformed = backend.concatenate(
             [combined_media_transformed, organic_rf_transformed], axis=-1
         )
-        combined_beta = tf.concat([combined_beta, beta_gorf], axis=-1)
+        combined_beta = backend.concatenate([combined_beta, beta_gorf], axis=-1)
 
-      sigma_gt = tf.transpose(tf.broadcast_to(sigma, [n_times, n_geos]))
-      y_pred_combined_media = tau_gt + tf.einsum(
+      sigma_gt = backend.transpose(
+          backend.broadcast_to(sigma, [n_times, n_geos])
+      )
+      y_pred_combined_media = tau_gt + backend.einsum(
           "gtm,gm->gt", combined_media_transformed, combined_beta
       )
       # Omit gamma_c, xi_c, and gamma_gc from joint distribution output if
@@ -387,22 +392,22 @@ class PosteriorMCMCSampler:
       if n_controls:
         gamma_c = yield prior_broadcast.gamma_c
         xi_c = yield prior_broadcast.xi_c
-        gamma_gc_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
+        gamma_gc_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_controls],
             name=constants.GAMMA_GC_DEV,
         )
-        gamma_gc = yield tfp.distributions.Deterministic(
+        gamma_gc = yield backend.tfd.Deterministic(
             gamma_c + xi_c * gamma_gc_dev, name=constants.GAMMA_GC
         )
-        y_pred_combined_media += tf.einsum(
+        y_pred_combined_media += backend.einsum(
             "gtc,gc->gt", controls_scaled, gamma_gc
         )
 
       if mmm.non_media_treatments is not None:
         xi_n = yield prior_broadcast.xi_n
-        gamma_gn_dev = yield tfp.distributions.Sample(
-            tfp.distributions.Normal(0, 1),
+        gamma_gn_dev = yield backend.tfd.Sample(
+            backend.tfd.Normal(0, 1),
             [n_geos, n_non_media_channels],
             name=constants.GAMMA_GN_DEV,
         )
@@ -425,15 +430,15 @@ class PosteriorMCMCSampler:
               eta_x=xi_n,
               beta_gx_dev=gamma_gn_dev,
           )
-          gamma_n = yield tfp.distributions.Deterministic(
+          gamma_n = yield backend.tfd.Deterministic(
               gamma_n_value, name=constants.GAMMA_N
           )
         else:
           raise ValueError(f"Unsupported prior type: {prior_type}")
-        gamma_gn = yield tfp.distributions.Deterministic(
+        gamma_gn = yield backend.tfd.Deterministic(
             gamma_n + xi_n * gamma_gn_dev, name=constants.GAMMA_GN
         )
-        y_pred = y_pred_combined_media + tf.einsum(
+        y_pred = y_pred_combined_media + backend.einsum(
             "gtn,gn->gt", non_media_treatments_normalized, gamma_gn
         )
       else:
@@ -445,21 +450,19 @@ class PosteriorMCMCSampler:
       # 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"
-        )
+        y_pred_holdout = backend.where(holdout_id, 0.0, y_pred)
+        test_sd = backend.cast(1.0 / np.sqrt(2.0 * np.pi), backend.float32)
+        sigma_gt_holdout = backend.where(holdout_id, test_sd, sigma_gt)
+        yield backend.tfd.Normal(y_pred_holdout, sigma_gt_holdout, name="y")
       else:
-        yield tfp.distributions.Normal(y_pred, sigma_gt, name="y")
+        yield backend.tfd.Normal(y_pred, sigma_gt, name="y")
 
     return joint_dist_unpinned
 
-  def _get_joint_dist(self) -> tfp.distributions.Distribution:
+  def _get_joint_dist(self) -> backend.tfd.Distribution:
     mmm = self.model
     y = (
-        tf.where(mmm.holdout_id, 0.0, mmm.kpi_scaled)
+        backend.where(mmm.holdout_id, 0.0, mmm.kpi_scaled)
         if mmm.holdout_id is not None
         else mmm.kpi_scaled
     )
@@ -471,7 +474,7 @@ class PosteriorMCMCSampler:
       n_adapt: int,
       n_burnin: int,
       n_keep: int,
-      current_state: Mapping[str, tf.Tensor] | None = None,
+      current_state: Mapping[str, backend.Tensor] | None = None,
       init_step_size: int | None = None,
       dual_averaging_kwargs: Mapping[str, int] | None = None,
       max_tree_depth: int = 10,
@@ -528,7 +531,7 @@ class PosteriorMCMCSampler:
         be a positive integer. For more information, see `tf.while_loop`.
       seed: An `int32[2]` Tensor or a Python list or tuple of 2 `int`s, which
         will be treated as stateless seeds; or a Python `int` or `None`, which
-        will be treated as stateful seeds. See [tfp.random.sanitize_seed]
+        will be converted into a stateless seed. See [tfp.random.sanitize_seed]
         (https://www.tensorflow.org/probability/api_docs/python/tfp/random/sanitize_seed).
       **pins: These are used to condition the provided joint distribution, and
         are passed directly to `joint_dist.experimental_pin(**pins)`.
@@ -547,7 +550,9 @@ class PosteriorMCMCSampler:
           " [tfp.random.sanitize_seed](https://www.tensorflow.org/probability/api_docs/python/tfp/random/sanitize_seed)"
           " for details."
       )
-    seed = tfp.random.sanitize_seed(seed) if seed is not None else None
+    if seed is not None and isinstance(seed, int):
+      seed = (seed, seed)
+    seed = backend.random.sanitize_seed(seed) if seed is not None else None
     n_chains_list = [n_chains] if isinstance(n_chains, int) else n_chains
     total_chains = np.sum(n_chains_list)
 
@@ -570,7 +575,7 @@ class PosteriorMCMCSampler:
             seed=seed,
             **pins,
         )
-      except tf.errors.ResourceExhaustedError as error:
+      except backend.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"
@@ -582,9 +587,11 @@ class PosteriorMCMCSampler:
       traces.append(mcmc.trace)
 
     mcmc_states = {
-        k: tf.einsum(
+        k: backend.einsum(
             "ij...->ji...",
-            tf.concat([state[k] for state in states], axis=1)[n_burnin:, ...],
+            backend.concatenate([state[k] for state in states], axis=1)[
+                n_burnin:, ...
+            ],
         )
         for k in states[0].keys()
         if k not in constants.UNSAVED_PARAMETERS
@@ -602,10 +609,10 @@ class PosteriorMCMCSampler:
     mcmc_trace = {}
     for k in traces[0].keys():
       if k not in constants.IGNORED_TRACE_METRICS:
-        mcmc_trace[k] = tf.concat(
+        mcmc_trace[k] = backend.concatenate(
             [
-                tf.broadcast_to(
-                    tf.transpose(trace[k][n_burnin:, ...]),
+                backend.broadcast_to(
+                    backend.transpose(trace[k][n_burnin:, ...]),
                     [n_chains_list[i], n_keep],
                 )
                 for i, trace in enumerate(traces)
@@ -645,7 +652,7 @@ class PosteriorMCMCSampler:
     # 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] = backend.broadcast_to(
         sample_stats[constants.STEP_SIZE], [total_chains, n_keep]
     )
     sample_stats_dims[constants.STEP_SIZE] = [constants.CHAIN, constants.DRAW]