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

meridian/model/prior_distribution.py

@@ -19,18 +19,21 @@ used by the Meridian model object.
 """
 
 from __future__ import annotations
-from collections.abc import MutableMapping
+from collections.abc import MutableMapping, Sequence
 import dataclasses
 from typing import Any
 import warnings
+
+from meridian import backend
 from meridian import constants
+
 import numpy as np
-import tensorflow as tf
-import tensorflow_probability as tfp
 
 
 __all__ = [
+    'IndependentMultivariateDistribution',
     'PriorDistribution',
+    'distributions_are_equal',
 ]
 
 
@@ -203,9 +206,9 @@ class PriorDistribution:
       distribution is `HalfNormal(5.0)`.
     roi_m: Prior distribution on the ROI of each media channel. This parameter
       is only used when `paid_media_prior_type` is `'roi'`, in which case
-      `beta_m` is calculated as a deterministic function of `roi_rf`,
-      `alpha_rf`, `ec_rf`, `slope_rf`, and the spend associated with each media
-      channel. Default distribution is `LogNormal(0.2, 0.9)`. When `kpi_type` is
+      `beta_m` is calculated as a deterministic function of `roi_m`, `alpha_m`,
+      `ec_m`, `slope_m`, and the spend associated with each media channel.
+      Default distribution is `LogNormal(0.2, 0.9)`. When `kpi_type` is
       `'non_revenue'` and `revenue_per_kpi` is not provided, ROI is interpreted
       as incremental KPI units per monetary unit spent. In this case, the
       default value for `roi_m` and `roi_rf` will be ignored and a common ROI
@@ -270,196 +273,202 @@ class PriorDistribution:
       `TruncatedNormal(0.0, 0.1, -1.0, 1.0)`.
   """
 
-  knot_values: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Normal(
+  knot_values: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Normal(
           0.0, 5.0, name=constants.KNOT_VALUES
       ),
   )
-  tau_g_excl_baseline: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Normal(
+  tau_g_excl_baseline: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Normal(
           0.0, 5.0, name=constants.TAU_G_EXCL_BASELINE
       ),
   )
-  beta_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  beta_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           5.0, name=constants.BETA_M
       ),
   )
-  beta_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  beta_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           5.0, name=constants.BETA_RF
       ),
   )
-  beta_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  beta_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           5.0, name=constants.BETA_OM
       ),
   )
-  beta_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  beta_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           5.0, name=constants.BETA_ORF
       ),
   )
-  eta_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
-          1.0, name=constants.ETA_M
-      ),
+  eta_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(1.0, name=constants.ETA_M),
   )
-  eta_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  eta_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           1.0, name=constants.ETA_RF
       ),
   )
-  eta_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  eta_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           1.0, name=constants.ETA_OM
       ),
   )
-  eta_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
+  eta_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(
           1.0, name=constants.ETA_ORF
       ),
   )
-  gamma_c: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Normal(
+  gamma_c: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Normal(
           0.0, 5.0, name=constants.GAMMA_C
       ),
   )
-  gamma_n: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Normal(
+  gamma_n: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Normal(
           0.0, 5.0, name=constants.GAMMA_N
       ),
   )
-  xi_c: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
-          5.0, name=constants.XI_C
-      ),
+  xi_c: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(5.0, name=constants.XI_C),
   )
-  xi_n: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
-          5.0, name=constants.XI_N
-      ),
+  xi_n: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(5.0, name=constants.XI_N),
   )
-  alpha_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Uniform(
+  alpha_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Uniform(
           0.0, 1.0, name=constants.ALPHA_M
       ),
   )
-  alpha_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Uniform(
+  alpha_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Uniform(
           0.0, 1.0, name=constants.ALPHA_RF
       ),
   )
-  alpha_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Uniform(
+  alpha_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Uniform(
           0.0, 1.0, name=constants.ALPHA_OM
       ),
   )
-  alpha_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Uniform(
+  alpha_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Uniform(
           0.0, 1.0, name=constants.ALPHA_ORF
       ),
   )
-  ec_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.TruncatedNormal(
+  ec_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.TruncatedNormal(
           0.8, 0.8, 0.1, 10, name=constants.EC_M
       ),
   )
-  ec_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.TransformedDistribution(
-          tfp.distributions.LogNormal(0.7, 0.4),
-          tfp.bijectors.Shift(0.1),
+  ec_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.TransformedDistribution(
+          backend.tfd.LogNormal(0.7, 0.4),
+          backend.bijectors.Shift(0.1),
           name=constants.EC_RF,
       ),
   )
-  ec_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.TruncatedNormal(
+  ec_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.TruncatedNormal(
           0.8, 0.8, 0.1, 10, name=constants.EC_OM
       ),
   )
-  ec_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.TransformedDistribution(
-          tfp.distributions.LogNormal(0.7, 0.4),
-          tfp.bijectors.Shift(0.1),
+  ec_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.TransformedDistribution(
+          backend.tfd.LogNormal(0.7, 0.4),
+          backend.bijectors.Shift(0.1),
           name=constants.EC_ORF,
       ),
   )
-  slope_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Deterministic(
+  slope_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Deterministic(
           1.0, name=constants.SLOPE_M
       ),
   )
-  slope_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  slope_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.7, 0.4, name=constants.SLOPE_RF
       ),
   )
-  slope_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Deterministic(
+  slope_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Deterministic(
           1.0, name=constants.SLOPE_OM
       ),
   )
-  slope_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  slope_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.7, 0.4, name=constants.SLOPE_ORF
       ),
   )
-  sigma: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.HalfNormal(
-          5.0, name=constants.SIGMA
-      ),
+  sigma: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.HalfNormal(5.0, name=constants.SIGMA),
   )
-  roi_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  roi_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.2, 0.9, name=constants.ROI_M
       ),
   )
-  roi_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  roi_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.2, 0.9, name=constants.ROI_RF
       ),
   )
-  mroi_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  mroi_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.0, 0.5, name=constants.MROI_M
       ),
   )
-  mroi_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.LogNormal(
+  mroi_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.LogNormal(
           0.0, 0.5, name=constants.MROI_RF
       ),
   )
-  contribution_m: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Beta(
+  contribution_m: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Beta(
           1.0, 99.0, name=constants.CONTRIBUTION_M
       ),
   )
-  contribution_rf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Beta(
+  contribution_rf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Beta(
           1.0, 99.0, name=constants.CONTRIBUTION_RF
       ),
   )
-  contribution_om: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Beta(
+  contribution_om: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Beta(
           1.0, 99.0, name=constants.CONTRIBUTION_OM
       ),
   )
-  contribution_orf: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.Beta(
+  contribution_orf: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.Beta(
           1.0, 99.0, name=constants.CONTRIBUTION_ORF
       ),
   )
-  contribution_n: tfp.distributions.Distribution = dataclasses.field(
-      default_factory=lambda: tfp.distributions.TruncatedNormal(
+  contribution_n: backend.tfd.Distribution = dataclasses.field(
+      default_factory=lambda: backend.tfd.TruncatedNormal(
           loc=0.0, scale=0.1, low=-1.0, high=1.0, name=constants.CONTRIBUTION_N
       ),
   )
 
+  def __post_init__(self):
+    for param, bounds in _parameter_space_bounds.items():
+      prevent_deterministic_prior_at_bounds = (
+          _prevent_deterministic_prior_at_bounds[param]
+          if param in _prevent_deterministic_prior_at_bounds.keys()
+          else (False, False)
+      )
+      _validate_support(
+          param,
+          getattr(self, param),
+          bounds,
+          prevent_deterministic_prior_at_bounds,
+      )
+
   def __setstate__(self, state):
     # Override to support pickling.
     def _unpack_distribution_params(
         params: MutableMapping[str, Any],
-    ) -> tfp.distributions.Distribution:
+    ) -> backend.tfd.Distribution:
       if constants.DISTRIBUTION in params:
         params[constants.DISTRIBUTION] = _unpack_distribution_params(
             params[constants.DISTRIBUTION]
@@ -478,7 +487,7 @@ class PriorDistribution:
     state = self.__dict__.copy()
 
     def _pack_distribution_params(
-        dist: tfp.distributions.Distribution,
+        dist: backend.tfd.Distribution,
     ) -> MutableMapping[str, Any]:
       params = dist.parameters
       params[constants.DISTRIBUTION_TYPE] = type(dist)
@@ -493,11 +502,9 @@ class PriorDistribution:
 
     return state
 
-  def has_deterministic_param(
-      self, param: tfp.distributions.Distribution
-  ) -> bool:
+  def has_deterministic_param(self, param: backend.tfd.Distribution) -> bool:
     return hasattr(self, param) and isinstance(
-        getattr(self, param).distribution, tfp.distributions.Deterministic
+        getattr(self, param).distribution, backend.tfd.Deterministic
     )
 
   def broadcast(
@@ -550,7 +557,7 @@ class PriorDistribution:
     """
 
     def _validate_media_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if (
           param.batch_shape.as_list()
@@ -573,7 +580,7 @@ class PriorDistribution:
     _validate_media_custom_priors(self.beta_m)
 
     def _validate_organic_media_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if (
           param.batch_shape.as_list()
@@ -595,7 +602,7 @@ class PriorDistribution:
     _validate_organic_media_custom_priors(self.beta_om)
 
     def _validate_organic_rf_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if (
           param.batch_shape.as_list()
@@ -617,7 +624,7 @@ class PriorDistribution:
     _validate_organic_rf_custom_priors(self.beta_orf)
 
     def _validate_rf_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if param.batch_shape.as_list() and n_rf_channels != param.batch_shape[0]:
         raise ValueError(
@@ -637,7 +644,7 @@ class PriorDistribution:
     _validate_rf_custom_priors(self.beta_rf)
 
     def _validate_control_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if param.batch_shape.as_list() and n_controls != param.batch_shape[0]:
         raise ValueError(
@@ -651,7 +658,7 @@ class PriorDistribution:
     _validate_control_custom_priors(self.xi_c)
 
     def _validate_non_media_custom_priors(
-        param: tfp.distributions.Distribution,
+        param: backend.tfd.Distribution,
     ) -> None:
       if (
           param.batch_shape.as_list()
@@ -669,7 +676,7 @@ class PriorDistribution:
     _validate_non_media_custom_priors(self.gamma_n)
     _validate_non_media_custom_priors(self.xi_n)
 
-    knot_values = tfp.distributions.BatchBroadcast(
+    knot_values = backend.tfd.BatchBroadcast(
         self.knot_values,
         n_knots,
         name=constants.KNOT_VALUES,
@@ -680,19 +687,19 @@ class PriorDistribution:
       )
     else:
       tau_g_converted = self.tau_g_excl_baseline
-    tau_g_excl_baseline = tfp.distributions.BatchBroadcast(
+    tau_g_excl_baseline = backend.tfd.BatchBroadcast(
         tau_g_converted, n_geos - 1, name=constants.TAU_G_EXCL_BASELINE
     )
-    beta_m = tfp.distributions.BatchBroadcast(
+    beta_m = backend.tfd.BatchBroadcast(
         self.beta_m, n_media_channels, name=constants.BETA_M
     )
-    beta_rf = tfp.distributions.BatchBroadcast(
+    beta_rf = backend.tfd.BatchBroadcast(
         self.beta_rf, n_rf_channels, name=constants.BETA_RF
     )
-    beta_om = tfp.distributions.BatchBroadcast(
+    beta_om = backend.tfd.BatchBroadcast(
         self.beta_om, n_organic_media_channels, name=constants.BETA_OM
     )
-    beta_orf = tfp.distributions.BatchBroadcast(
+    beta_orf = backend.tfd.BatchBroadcast(
         self.beta_orf, n_organic_rf_channels, name=constants.BETA_ORF
     )
     if is_national:
@@ -705,66 +712,66 @@ class PriorDistribution:
       eta_rf_converted = self.eta_rf
       eta_om_converted = self.eta_om
       eta_orf_converted = self.eta_orf
-    eta_m = tfp.distributions.BatchBroadcast(
+    eta_m = backend.tfd.BatchBroadcast(
         eta_m_converted, n_media_channels, name=constants.ETA_M
     )
-    eta_rf = tfp.distributions.BatchBroadcast(
+    eta_rf = backend.tfd.BatchBroadcast(
         eta_rf_converted, n_rf_channels, name=constants.ETA_RF
     )
-    eta_om = tfp.distributions.BatchBroadcast(
+    eta_om = backend.tfd.BatchBroadcast(
         eta_om_converted,
         n_organic_media_channels,
         name=constants.ETA_OM,
     )
-    eta_orf = tfp.distributions.BatchBroadcast(
+    eta_orf = backend.tfd.BatchBroadcast(
         eta_orf_converted, n_organic_rf_channels, name=constants.ETA_ORF
     )
-    gamma_c = tfp.distributions.BatchBroadcast(
+    gamma_c = backend.tfd.BatchBroadcast(
         self.gamma_c, n_controls, name=constants.GAMMA_C
     )
     if is_national:
       xi_c_converted = _convert_to_deterministic_0_distribution(self.xi_c)
     else:
       xi_c_converted = self.xi_c
-    xi_c = tfp.distributions.BatchBroadcast(
+    xi_c = backend.tfd.BatchBroadcast(
         xi_c_converted, n_controls, name=constants.XI_C
     )
-    gamma_n = tfp.distributions.BatchBroadcast(
+    gamma_n = backend.tfd.BatchBroadcast(
         self.gamma_n, n_non_media_channels, name=constants.GAMMA_N
     )
     if is_national:
       xi_n_converted = _convert_to_deterministic_0_distribution(self.xi_n)
     else:
       xi_n_converted = self.xi_n
-    xi_n = tfp.distributions.BatchBroadcast(
+    xi_n = backend.tfd.BatchBroadcast(
         xi_n_converted, n_non_media_channels, name=constants.XI_N
     )
-    alpha_m = tfp.distributions.BatchBroadcast(
+    alpha_m = backend.tfd.BatchBroadcast(
         self.alpha_m, n_media_channels, name=constants.ALPHA_M
     )
-    alpha_rf = tfp.distributions.BatchBroadcast(
+    alpha_rf = backend.tfd.BatchBroadcast(
         self.alpha_rf, n_rf_channels, name=constants.ALPHA_RF
     )
-    alpha_om = tfp.distributions.BatchBroadcast(
+    alpha_om = backend.tfd.BatchBroadcast(
         self.alpha_om, n_organic_media_channels, name=constants.ALPHA_OM
     )
-    alpha_orf = tfp.distributions.BatchBroadcast(
+    alpha_orf = backend.tfd.BatchBroadcast(
         self.alpha_orf, n_organic_rf_channels, name=constants.ALPHA_ORF
     )
-    ec_m = tfp.distributions.BatchBroadcast(
+    ec_m = backend.tfd.BatchBroadcast(
         self.ec_m, n_media_channels, name=constants.EC_M
     )
-    ec_rf = tfp.distributions.BatchBroadcast(
+    ec_rf = backend.tfd.BatchBroadcast(
         self.ec_rf, n_rf_channels, name=constants.EC_RF
     )
-    ec_om = tfp.distributions.BatchBroadcast(
+    ec_om = backend.tfd.BatchBroadcast(
         self.ec_om, n_organic_media_channels, name=constants.EC_OM
     )
-    ec_orf = tfp.distributions.BatchBroadcast(
+    ec_orf = backend.tfd.BatchBroadcast(
         self.ec_orf, n_organic_rf_channels, name=constants.EC_ORF
     )
     if (
-        not isinstance(self.slope_m, tfp.distributions.Deterministic)
+        not isinstance(self.slope_m, backend.tfd.Deterministic)
         or (np.isscalar(self.slope_m.loc.numpy()) and self.slope_m.loc != 1.0)
         or (
             self.slope_m.batch_shape.as_list()
@@ -776,14 +783,14 @@ class PriorDistribution:
           ' This may lead to poor MCMC convergence and budget optimization'
           ' may no longer produce a global optimum.'
       )
-    slope_m = tfp.distributions.BatchBroadcast(
+    slope_m = backend.tfd.BatchBroadcast(
         self.slope_m, n_media_channels, name=constants.SLOPE_M
     )
-    slope_rf = tfp.distributions.BatchBroadcast(
+    slope_rf = backend.tfd.BatchBroadcast(
         self.slope_rf, n_rf_channels, name=constants.SLOPE_RF
     )
     if (
-        not isinstance(self.slope_om, tfp.distributions.Deterministic)
+        not isinstance(self.slope_om, backend.tfd.Deterministic)
         or (np.isscalar(self.slope_om.loc.numpy()) and self.slope_om.loc != 1.0)
         or (
             self.slope_om.batch_shape.as_list()
@@ -795,16 +802,16 @@ class PriorDistribution:
           ' This may lead to poor MCMC convergence and budget optimization'
           ' may no longer produce a global optimum.'
       )
-    slope_om = tfp.distributions.BatchBroadcast(
+    slope_om = backend.tfd.BatchBroadcast(
         self.slope_om, n_organic_media_channels, name=constants.SLOPE_OM
     )
-    slope_orf = tfp.distributions.BatchBroadcast(
+    slope_orf = backend.tfd.BatchBroadcast(
         self.slope_orf, n_organic_rf_channels, name=constants.SLOPE_ORF
     )
 
     # If `unique_sigma_for_each_geo == False`, then make a scalar batch.
     sigma_shape = n_geos if (n_geos > 1 and unique_sigma_for_each_geo) else []
-    sigma = tfp.distributions.BatchBroadcast(
+    sigma = backend.tfd.BatchBroadcast(
         self.sigma, sigma_shape, name=constants.SIGMA
     )
 
@@ -818,37 +825,37 @@ class PriorDistribution:
     else:
       roi_m_converted = self.roi_m
       roi_rf_converted = self.roi_rf
-    roi_m = tfp.distributions.BatchBroadcast(
+    roi_m = backend.tfd.BatchBroadcast(
         roi_m_converted, n_media_channels, name=constants.ROI_M
     )
-    roi_rf = tfp.distributions.BatchBroadcast(
+    roi_rf = backend.tfd.BatchBroadcast(
         roi_rf_converted, n_rf_channels, name=constants.ROI_RF
     )
 
-    mroi_m = tfp.distributions.BatchBroadcast(
+    mroi_m = backend.tfd.BatchBroadcast(
         self.mroi_m, n_media_channels, name=constants.MROI_M
     )
-    mroi_rf = tfp.distributions.BatchBroadcast(
+    mroi_rf = backend.tfd.BatchBroadcast(
         self.mroi_rf, n_rf_channels, name=constants.MROI_RF
     )
 
-    contribution_m = tfp.distributions.BatchBroadcast(
+    contribution_m = backend.tfd.BatchBroadcast(
         self.contribution_m, n_media_channels, name=constants.CONTRIBUTION_M
     )
-    contribution_rf = tfp.distributions.BatchBroadcast(
+    contribution_rf = backend.tfd.BatchBroadcast(
         self.contribution_rf, n_rf_channels, name=constants.CONTRIBUTION_RF
     )
-    contribution_om = tfp.distributions.BatchBroadcast(
+    contribution_om = backend.tfd.BatchBroadcast(
         self.contribution_om,
         n_organic_media_channels,
         name=constants.CONTRIBUTION_OM,
     )
-    contribution_orf = tfp.distributions.BatchBroadcast(
+    contribution_orf = backend.tfd.BatchBroadcast(
         self.contribution_orf,
         n_organic_rf_channels,
         name=constants.CONTRIBUTION_ORF,
     )
-    contribution_n = tfp.distributions.BatchBroadcast(
+    contribution_n = backend.tfd.BatchBroadcast(
         self.contribution_n, n_non_media_channels, name=constants.CONTRIBUTION_N
     )
 
@@ -892,9 +899,283 @@ class PriorDistribution:
     )
 
 
+class IndependentMultivariateDistribution(backend.tfd.Distribution):
+  """Container for a joint distribution created from independent distributions.
+
+  This class is useful when one wants to define a joint distribution for a
+  Meridian prior, where the elements are not necessarily from the same
+  distribution family. For example, to define a distribution where
+  one element is Uniform and the second is triangular:
+
+  ```python
+  distributions = [
+      tfp.distributions.Uniform(0.0, 1.0),
+      tfp.distributions.Triangular(0.0, 1.0, 0.5)
+      ]
+  distribution = IndependentMultivariateDistribution(distributions)
+  ```
+
+  It is also possible to define a distribution where multiple elements come
+  from the same distribution family. For example, to define a distribution where
+  the three elements are LogNormal(0.2, 0.9), LogNormal(0, 0.5) and
+  Gamma(2, 2):
+
+  ```python
+  distributions = [
+      tfp.distributions.LogNormal([0.2, 0.0], [0.9, 0.5]),
+      tfp.distributions.Gamma(2.0, 2.0)
+      ]
+  distribution = IndependentMultivariateDistribution(distributions)
+  ```
+
+  This class cannot contain instances of `tfd.Deterministic`.
+  """
+
+  def __init__(
+      self,
+      distributions: Sequence[backend.tfd.Distribution],
+      validate_args: bool = False,
+      allow_nan_stats: bool = True,
+      name: str | None = None,
+  ):
+    """Initializes a batch of independent distributions from different families.
+
+    Args:
+      distributions: List of `tfd.Distribution` from which to construct a
+        multivariate distribution. The distributions must have scalar or one
+        dimensional batch shapes; the resulting batch shape will be the sum of
+        the underlying batch shapes.
+      validate_args: Python `bool`. When `True` distribution parameters are
+        checked for validity despite possibly degrading runtime performance.
+        When `False` invalid inputs may silently render incorrect outputs.
+        Default value is `False`.
+      allow_nan_stats: Python `bool`. When `True`, statistics (e.g., mean, mode,
+        variance) use the value "`NaN`" to indicate the result is undefined.
+        When `False`, an exception is raised if one or more of the statistic's
+        batch members are undefined. Default value is `True`.
+      name: Python `str` name prefixed to Ops created by this class. Default
+        value is 'IndependentMultivariate' followed by the names of the
+        underlying distributions.
+
+    Raises:
+        ValueError: If one or more distributions are instances of
+        `tfd.Deterministic` or dtypes differ between the
+        distributions.
+    """
+    parameters = dict(locals())
+
+    self._verify_distributions(distributions)
+
+    self._distributions = [
+        dist
+        if not dist.is_scalar_batch()
+        else backend.tfd.BatchBroadcast(dist, (1,))
+        for dist in distributions
+    ]
+
+    self._distribution_batch_shapes = self._get_distribution_batch_shapes()
+    self._distribution_batch_shape_tensors = backend.concatenate(
+        [dist.batch_shape_tensor() for dist in self._distributions],
+        axis=0,
+    )
+
+    dtype = self._verify_dtypes()
+
+    name = name or '-'.join(
+        [constants.INDEPENDENT_MULTIVARIATE] + [d.name for d in distributions]
+    )
+
+    super().__init__(
+        dtype=dtype,
+        reparameterization_type=backend.tfd.NOT_REPARAMETERIZED,
+        validate_args=validate_args,
+        allow_nan_stats=allow_nan_stats,
+        parameters=parameters,
+        name=name,
+    )
+
+  def _verify_distributions(
+      self, distributions: Sequence[backend.tfd.Distribution]
+  ):
+    """Check for deterministic distributions and raise an error if found."""
+
+    if any(
+        isinstance(dist, backend.tfd.Deterministic)
+        for dist in distributions
+    ):
+      raise ValueError(
+          f'{self.__class__.__name__} cannot contain `Deterministic` '
+          'distributions. To implement a nearly deterministic element of this '
+          'distribution, we recommend using `backend.tfd.Uniform` with a '
+          'small range. For example to define a distribution that is nearly '
+          '`Deterministic(1.0)`, use '
+          '`tfp.distribution.Uniform(1.0 - 1e-9, 1.0 + 1e-9)`'
+      )
+
+  def _verify_dtypes(self) -> str:
+    dtypes = [dist.dtype for dist in self._distributions]
+    if len(set(dtypes)) != 1:
+      raise ValueError(
+          f'All distributions must have the same dtype. Found: {dtypes}.'
+      )
+
+    return backend.result_type(*dtypes)
+
+  def _event_shape(self):
+    return backend.TensorShape([])
+
+  def _batch_shape_tensor(self):
+    distribution_batch_shape_tensors = backend.concatenate(
+        [dist.batch_shape_tensor() for dist in self._distributions],
+        axis=0,
+    )
+    return backend.reduce_sum(
+        distribution_batch_shape_tensors, keepdims=True
+    )
+
+  def _batch_shape(self):
+    return backend.TensorShape(sum(self._distribution_batch_shapes))
+
+  def _sample_n(self, n, seed=None):
+    return backend.concatenate(
+        [dist.sample(n, seed) for dist in self._distributions], axis=-1
+    )
+
+  def _quantile(self, value):
+    value = self._broadcast_value(value)
+    split_value = backend.split(
+        value,
+        self._distribution_batch_shapes, axis=-1
+        )
+    quantiles = [
+        dist.quantile(sv) for dist, sv in zip(self._distributions, split_value)
+    ]
+
+    return backend.concatenate(quantiles, axis=-1)
+
+  def _log_prob(self, value):
+    value = self._broadcast_value(value)
+    split_value = backend.split(
+        value,
+        self._distribution_batch_shapes,
+        axis=-1
+        )
+    log_probs = [
+        dist.log_prob(sv) for dist, sv in zip(self._distributions, split_value)
+    ]
+
+    return backend.concatenate(log_probs, axis=-1)
+
+  def _log_cdf(self, value):
+    value = self._broadcast_value(value)
+    split_value = backend.split(
+        value,
+        self._distribution_batch_shapes,
+        axis=-1
+        )
+
+    log_cdfs = [
+        dist.log_cdf(sv) for dist, sv in zip(self._distributions, split_value)
+    ]
+
+    return backend.concatenate(log_cdfs, axis=-1)
+
+  def _mean(self):
+    return backend.concatenate(
+        [dist.mean() for dist in self._distributions], axis=0
+    )
+
+  def _variance(self):
+    return backend.concatenate(
+        [dist.variance() for dist in self._distributions], axis=0
+    )
+
+  def _default_event_space_bijector(self):
+    """Mapping from R^n to the event space of the wrapped distributions.
+
+    This is the blockwise concatenation of the underlying bijectors.
+
+    Returns:
+      A `tfp.bijectors.Blockwise` object that concatenates the underlying
+      bijectors.
+    """
+    bijectors = [
+        d.experimental_default_event_space_bijector()
+        for d in self._distributions
+    ]
+
+    return backend.bijectors.Blockwise(
+        bijectors,
+        block_sizes=self._distribution_batch_shapes,
+    )
+
+  def _broadcast_value(self, value: backend.Tensor) -> backend.Tensor:
+    value = backend.to_tensor(value)
+    broadcast_shape = backend.broadcast_dynamic_shape(
+        value.shape, self.batch_shape_tensor()
+    )
+    return backend.broadcast_to(value, broadcast_shape)
+
+  def _get_distribution_batch_shapes(self) -> Sequence[int]:
+    """Sequence of batch shapes of underlying distributions."""
+
+    batch_shapes = []
+
+    for dist in self._distributions:
+      try:
+        (dist_batch_shape,) = dist.batch_shape
+      except ValueError as exc:
+        raise ValueError(
+            'All distributions must be 0- or 1-dimensional.'
+            f' Found {len(dist.batch_shape)}-dimensional distribution:'
+            f' {dist.batch_shape}.'
+        ) from exc
+      else:
+        batch_shapes.append(dist_batch_shape)
+
+    return batch_shapes
+
+
+def distributions_are_equal(
+    a: backend.tfd.Distribution, b: backend.tfd.Distribution
+) -> bool:
+  """Determine if two distributions are equal."""
+  if type(a) != type(b):  # pylint: disable=unidiomatic-typecheck
+    return False
+
+  a_params = a.parameters.copy()
+  b_params = b.parameters.copy()
+
+  if constants.DISTRIBUTION in a_params and constants.DISTRIBUTION in b_params:
+    if not distributions_are_equal(
+        a_params[constants.DISTRIBUTION], b_params[constants.DISTRIBUTION]
+    ):
+      return False
+    del a_params[constants.DISTRIBUTION]
+    del b_params[constants.DISTRIBUTION]
+
+  if constants.DISTRIBUTION in a_params or constants.DISTRIBUTION in b_params:
+    return False
+
+  if a_params.keys() != b_params.keys():
+    return False
+
+  for key in a_params.keys():
+    if isinstance(
+        a_params[key], (backend.Tensor, np.ndarray, float, int)
+    ) and isinstance(b_params[key], (backend.Tensor, np.ndarray, float, int)):
+      if not backend.allclose(a_params[key], b_params[key]):
+        return False
+    else:
+      if a_params[key] != b_params[key]:
+        return False
+
+  return True
+
+
 def _convert_to_deterministic_0_distribution(
-    distribution: tfp.distributions.Distribution,
-) -> tfp.distributions.Distribution:
+    distribution: backend.tfd.Distribution,
+) -> backend.tfd.Distribution:
   """Converts the given distribution to a `Deterministic(0)` one.
 
   Args:
@@ -909,7 +1190,7 @@ def _convert_to_deterministic_0_distribution(
     distribution.
   """
   if (
-      not isinstance(distribution, tfp.distributions.Deterministic)
+      not isinstance(distribution, backend.tfd.Deterministic)
       or distribution.loc != 0
   ):
     warnings.warn(
@@ -917,7 +1198,7 @@ def _convert_to_deterministic_0_distribution(
         f' for national models. {distribution.name} has been automatically set'
         ' to Deterministic(0).'
     )
-    return tfp.distributions.Deterministic(loc=0, name=distribution.name)
+    return backend.tfd.Deterministic(loc=0, name=distribution.name)
   else:
     return distribution
 
@@ -928,7 +1209,7 @@ def _get_total_media_contribution_prior(
     name: str,
     p_mean: float = constants.P_MEAN,
     p_sd: float = constants.P_SD,
-) -> tfp.distributions.Distribution:
+) -> backend.tfd.Distribution:
   """Determines ROI priors based on total media contribution.
 
   Args:
@@ -945,34 +1226,123 @@ def _get_total_media_contribution_prior(
   """
   roi_mean = p_mean * kpi / np.sum(total_spend)
   roi_sd = p_sd * kpi / np.sqrt(np.sum(np.power(total_spend, 2)))
-  lognormal_sigma = tf.cast(
-      np.sqrt(np.log(roi_sd**2 / roi_mean**2 + 1)), dtype=tf.float32
+  lognormal_sigma = backend.cast(
+      np.sqrt(np.log(roi_sd**2 / roi_mean**2 + 1)), dtype=backend.float32
   )
-  lognormal_mu = tf.cast(
-      np.log(roi_mean * np.exp(-(lognormal_sigma**2) / 2)), dtype=tf.float32
+  lognormal_mu = backend.cast(
+      np.log(roi_mean * np.exp(-(lognormal_sigma**2) / 2)),
+      dtype=backend.float32,
   )
-  return tfp.distributions.LogNormal(lognormal_mu, lognormal_sigma, name=name)
+  return backend.tfd.LogNormal(lognormal_mu, lognormal_sigma, name=name)
 
 
-def distributions_are_equal(
-    a: tfp.distributions.Distribution, b: tfp.distributions.Distribution
-) -> bool:
-  """Determine if two distributions are equal."""
-  if type(a) != type(b):  # pylint: disable=unidiomatic-typecheck
-    return False
+def _validate_support(
+    parameter_name: str,
+    tfp_dist: backend.tfp.distributions.Distribution,
+    bounds: tuple[float, float],
+    prevent_deterministic_prior_at_bounds: tuple[bool, bool],
+) -> None:
+  """Validates that distribution support is within the parameter bounds.
 
-  a_params = a.parameters.copy()
-  b_params = b.parameters.copy()
+  Args:
+    parameter_name: Name of the parameter.
+    tfp_dist: The TFP distribution to validate.
+    bounds: Tuple containing the min and max values of the parameteter space.
+    prevent_deterministic_prior_at_bounds: Tuple of two booleans indicating
+      whether a deterministic prior is allowed at the lower and upper bounds,
+      respectively.
 
-  if constants.DISTRIBUTION in a_params and constants.DISTRIBUTION in b_params:
-    if not distributions_are_equal(
-        a_params[constants.DISTRIBUTION], b_params[constants.DISTRIBUTION]
-    ):
-      return False
-    del a_params[constants.DISTRIBUTION]
-    del b_params[constants.DISTRIBUTION]
+  Raises:
+    ValueError: If the distribution support is not within the parameter bounds.
+  """
+  # Note that `tfp.distributions.BatchBroadcast` objects have a `distribution`
+  # attribute that points to a `tfp.distributions.Distribution` object.
+  if isinstance(tfp_dist, backend.tfp.distributions.BatchBroadcast):
+    tfp_dist = tfp_dist.distribution
+  # Note that `tfp.distributions.Deterministic` does not have a `quantile`
+  # method implemented, so the min and max values must be extracted from the
+  # `loc` attribute instead.
+  if isinstance(
+      tfp_dist,
+      backend.tfp.python.distributions.deterministic.Deterministic
+  ):
+    support_min_vals = tfp_dist.loc
+    support_max_vals = tfp_dist.loc
+    for i in (0, 1):
+      if (
+          prevent_deterministic_prior_at_bounds[i]
+          and np.any(tfp_dist.loc == bounds[i])
+      ):
+        raise ValueError(
+            f'{parameter_name} was assigned a point mass (deterministic) prior'
+            f' at {bounds[i]}, which is not allowed.'
+        )
+  else:
+    try:
+      support_min_vals = tfp_dist.quantile(0)
+      support_max_vals = tfp_dist.quantile(1)
+    except (AttributeError, NotImplementedError):
+      warnings.warn(
+          f'The prior distribution for {parameter_name} does not have a'
+          f' `quantile` method implemented, so the support range validation'
+          f' was skipped. Confirm that your prior for {parameter_name} is'
+          f' appropriate.'
+      )
+      return
+  if np.any(support_min_vals < bounds[0]):
+    raise ValueError(
+        f'{parameter_name} was assigned a prior distribution that allows values'
+        f' less than the parameter minimum {bounds[0]}.'
+    )
+  if np.any(support_max_vals > bounds[1]):
+    raise ValueError(
+        f'{parameter_name} was assigned a prior distribution that allows values'
+        f' greater than the parameter maximum {bounds[1]}.'
+    )
 
-  if constants.DISTRIBUTION in a_params or constants.DISTRIBUTION in b_params:
-    return False
+# Dictionary of parameters that have a limited parameters space. The tuple
+# contains the lower and upper bounds, respectively.
+_parameter_space_bounds = {
+    'eta_m': (0, np.inf),
+    'eta_rf': (0, np.inf),
+    'eta_om': (0, np.inf),
+    'eta_orf': (0, np.inf),
+    'xi_c': (0, np.inf),
+    'xi_n': (0, np.inf),
+    'alpha_m': (0, 1),
+    'alpha_rf': (0, 1),
+    'alpha_om': (0, 1),
+    'alpha_orf': (0, 1),
+    'ec_m': (0, np.inf),
+    'ec_rf': (0, np.inf),
+    'ec_om': (0, np.inf),
+    'ec_orf': (0, np.inf),
+    'slope_m': (0, np.inf),
+    'slope_rf': (0, np.inf),
+    'slope_om': (0, np.inf),
+    'slope_orf': (0, np.inf),
+    'sigma': (0, np.inf),
+}
 
-  return a_params == b_params
+# Dictionary of parameters that do not allow a deterministic prior at one or
+# more of the parameter space bounds. The boolean tuple indicates whether a
+# deterministic prior is allowed at the lower bound or upper bound,
+# respectively, where `True` means "not allowed". This check is specifically for
+# point mass at finite paramteter space bounds, since point mass at infinity is
+# generally problematic for all parameters. Note that `sigma` should generally
+# not have point mass at zero, but this is not checked here because unit tests
+# require the ability to simulate data with `sigma` set to zero.
+_prevent_deterministic_prior_at_bounds = {
+    'alpha_m': (False, True),
+    'alpha_rf': (False, True),
+    'alpha_om': (False, True),
+    'alpha_orf': (False, True),
+    'ec_m': (True, False),
+    'ec_rf': (True, False),
+    'ec_om': (True, False),
+    'ec_orf': (True, False),
+    'slope_m': (True, False),
+    'slope_rf': (True, False),
+    'slope_om': (True, False),
+    'slope_orf': (True, False),
+}