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

meridian/model/model.py

@@ -22,6 +22,7 @@ import warnings
 
 import arviz as az
 import joblib
+from meridian import backend
 from meridian import constants
 from meridian.data import input_data as data
 from meridian.data import time_coordinates as tc
@@ -34,8 +35,6 @@ 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__ = [
@@ -70,7 +69,7 @@ def _warn_setting_national_args(**kwargs):
 
 
 def _check_for_negative_effect(
-    dist: tfp.distributions.Distribution, media_effects_dist: str
+    dist: backend.tfd.Distribution, media_effects_dist: str
 ):
   """Checks for negative effect in the model."""
   if (
@@ -202,45 +201,45 @@ class Meridian:
     return media.build_organic_rf_tensors(self.input_data)
 
   @functools.cached_property
-  def kpi(self) -> tf.Tensor:
-    return tf.convert_to_tensor(self.input_data.kpi, dtype=tf.float32)
+  def kpi(self) -> backend.Tensor:
+    return backend.to_tensor(self.input_data.kpi, dtype=backend.float32)
 
   @functools.cached_property
-  def revenue_per_kpi(self) -> tf.Tensor | None:
+  def revenue_per_kpi(self) -> backend.Tensor | None:
     if self.input_data.revenue_per_kpi is None:
       return None
-    return tf.convert_to_tensor(
-        self.input_data.revenue_per_kpi, dtype=tf.float32
+    return backend.to_tensor(
+        self.input_data.revenue_per_kpi, dtype=backend.float32
     )
 
   @functools.cached_property
-  def controls(self) -> tf.Tensor | None:
+  def controls(self) -> backend.Tensor | None:
     if self.input_data.controls is None:
       return None
-    return tf.convert_to_tensor(self.input_data.controls, dtype=tf.float32)
+    return backend.to_tensor(self.input_data.controls, dtype=backend.float32)
 
   @functools.cached_property
-  def non_media_treatments(self) -> tf.Tensor | None:
+  def non_media_treatments(self) -> backend.Tensor | None:
     if self.input_data.non_media_treatments is None:
       return None
-    return tf.convert_to_tensor(
-        self.input_data.non_media_treatments, dtype=tf.float32
+    return backend.to_tensor(
+        self.input_data.non_media_treatments, dtype=backend.float32
     )
 
   @functools.cached_property
-  def population(self) -> tf.Tensor:
-    return tf.convert_to_tensor(self.input_data.population, dtype=tf.float32)
+  def population(self) -> backend.Tensor:
+    return backend.to_tensor(self.input_data.population, dtype=backend.float32)
 
   @functools.cached_property
-  def total_spend(self) -> tf.Tensor:
-    return tf.convert_to_tensor(
-        self.input_data.get_total_spend(), dtype=tf.float32
+  def total_spend(self) -> backend.Tensor:
+    return backend.to_tensor(
+        self.input_data.get_total_spend(), dtype=backend.float32
     )
 
   @functools.cached_property
-  def total_outcome(self) -> tf.Tensor:
-    return tf.convert_to_tensor(
-        self.input_data.get_total_outcome(), dtype=tf.float32
+  def total_outcome(self) -> backend.Tensor:
+    return backend.to_tensor(
+        self.input_data.get_total_outcome(), dtype=backend.float32
     )
 
   @property
@@ -300,6 +299,8 @@ class Meridian:
     return knots.get_knot_info(
         n_times=self.n_times,
         knots=self.model_spec.knots,
+        enable_aks=self.model_spec.enable_aks,
+        data=self.input_data,
         is_national=self.is_national,
     )
 
@@ -312,8 +313,8 @@ class Meridian:
       return None
 
     if self.model_spec.control_population_scaling_id is not None:
-      controls_population_scaling_id = tf.convert_to_tensor(
-          self.model_spec.control_population_scaling_id, dtype=bool
+      controls_population_scaling_id = backend.to_tensor(
+          self.model_spec.control_population_scaling_id, dtype=backend.bool_
       )
     else:
       controls_population_scaling_id = None
@@ -332,8 +333,8 @@ class Meridian:
     if self.non_media_treatments is None:
       return None
     if self.model_spec.non_media_population_scaling_id is not None:
-      non_media_population_scaling_id = tf.convert_to_tensor(
-          self.model_spec.non_media_population_scaling_id, dtype=bool
+      non_media_population_scaling_id = backend.to_tensor(
+          self.model_spec.non_media_population_scaling_id, dtype=backend.bool_
       )
     else:
       non_media_population_scaling_id = None
@@ -349,7 +350,7 @@ class Meridian:
     return transformers.KpiTransformer(self.kpi, self.population)
 
   @functools.cached_property
-  def controls_scaled(self) -> tf.Tensor | None:
+  def controls_scaled(self) -> backend.Tensor | None:
     if self.controls is not None:
       # If `controls` is defined, then `controls_transformer` is also defined.
       return self.controls_transformer.forward(self.controls)  # pytype: disable=attribute-error
@@ -357,7 +358,7 @@ class Meridian:
       return None
 
   @functools.cached_property
-  def non_media_treatments_normalized(self) -> tf.Tensor | None:
+  def non_media_treatments_normalized(self) -> backend.Tensor | None:
     """Normalized non-media treatments.
 
     The non-media treatments values are scaled by population (for channels where
@@ -372,7 +373,7 @@ class Meridian:
       return None
 
   @functools.cached_property
-  def kpi_scaled(self) -> tf.Tensor:
+  def kpi_scaled(self) -> backend.Tensor:
     return self.kpi_transformer.forward(self.kpi)
 
   @functools.cached_property
@@ -416,14 +417,35 @@ class Meridian:
       # Geos are unique, so index is a 1-element array.
       return index[0]
     else:
-      return tf.argmax(self.population)
+      return backend.argmax(self.population)
 
   @functools.cached_property
-  def holdout_id(self) -> tf.Tensor | None:
+  def holdout_id(self) -> backend.Tensor | None:
     if self.model_spec.holdout_id is None:
       return None
-    tensor = tf.convert_to_tensor(self.model_spec.holdout_id, dtype=bool)
-    return tensor[tf.newaxis, ...] if self.is_national else tensor
+    tensor = backend.to_tensor(self.model_spec.holdout_id, dtype=backend.bool_)
+    return tensor[backend.newaxis, ...] if self.is_national else tensor
+
+  @functools.cached_property
+  def adstock_decay_spec(self) -> adstock_hill.AdstockDecaySpec:
+    """Returns `AdstockDecaySpec` object with correctly mapped channels."""
+    if isinstance(self.model_spec.adstock_decay_spec, str):
+      return adstock_hill.AdstockDecaySpec.from_consistent_type(
+          self.model_spec.adstock_decay_spec
+      )
+
+    try:
+      return self._create_adstock_decay_functions_from_channel_map(
+          self.model_spec.adstock_decay_spec
+      )
+    except KeyError as e:
+      raise ValueError(
+          "Unrecognized channel names found in `adstock_decay_spec` keys"
+          f" {tuple(self.model_spec.adstock_decay_spec.keys())}. Keys should"
+          " either contain only channel_names"
+          f" {tuple(self.input_data.get_all_adstock_hill_channels().tolist())} or"
+          " be one or more of {'media', 'rf', 'organic_media', 'organic_rf'}."
+      ) from e
 
   @functools.cached_property
   def prior_broadcast(self) -> prior_distribution.PriorDistribution:
@@ -469,7 +491,7 @@ class Meridian:
   def compute_non_media_treatments_baseline(
       self,
       non_media_baseline_values: Sequence[str | float] | None = None,
-  ) -> tf.Tensor:
+  ) -> backend.Tensor:
     """Computes the baseline for each non-media treatment channel.
 
     Args:
@@ -491,16 +513,19 @@ class Meridian:
     if non_media_baseline_values is None:
       non_media_baseline_values = self.model_spec.non_media_baseline_values
 
+    no_op_scaling_factor = backend.ones_like(self.population)[
+        :, backend.newaxis, backend.newaxis
+    ]
     if self.model_spec.non_media_population_scaling_id is not None:
-      scaling_factors = tf.where(
+      scaling_factors = backend.where(
           self.model_spec.non_media_population_scaling_id,
-          self.population[:, tf.newaxis, tf.newaxis],
-          tf.ones_like(self.population)[:, tf.newaxis, tf.newaxis],
+          self.population[:, backend.newaxis, backend.newaxis],
+          no_op_scaling_factor,
       )
     else:
-      scaling_factors = tf.ones_like(self.population)[:, tf.newaxis, tf.newaxis]
+      scaling_factors = no_op_scaling_factor
 
-    non_media_treatments_population_scaled = tf.math.divide_no_nan(
+    non_media_treatments_population_scaled = backend.divide_no_nan(
         self.non_media_treatments, scaling_factors
     )
 
@@ -528,15 +553,15 @@ class Meridian:
       baseline_value = non_media_baseline_values_filled[channel]
 
       if baseline_value == constants.NON_MEDIA_BASELINE_MIN:
-        baseline_for_channel = tf.reduce_min(
+        baseline_for_channel = backend.reduce_min(
             non_media_treatments_population_scaled[..., channel], axis=[0, 1]
         )
       elif baseline_value == constants.NON_MEDIA_BASELINE_MAX:
-        baseline_for_channel = tf.reduce_max(
+        baseline_for_channel = backend.reduce_max(
             non_media_treatments_population_scaled[..., channel], axis=[0, 1]
         )
       elif isinstance(baseline_value, numbers.Number):
-        baseline_for_channel = tf.cast(baseline_value, tf.float32)
+        baseline_for_channel = backend.cast(baseline_value, backend.float32)
       else:
         raise ValueError(
             f"Invalid non_media_baseline_values value: '{baseline_value}'. Only"
@@ -545,7 +570,7 @@ class Meridian:
 
       baseline_list.append(baseline_for_channel)
 
-    return tf.stack(baseline_list, axis=-1)
+    return backend.stack(baseline_list, axis=-1)
 
   def expand_selected_time_dims(
       self,
@@ -766,6 +791,58 @@ class Meridian:
           f" ({self.n_non_media_channels},)`."
       )
 
+  def _create_adstock_decay_functions_from_channel_map(
+      self, channel_function_map: Mapping[str, str]
+  ) -> adstock_hill.AdstockDecaySpec:
+    """Create `AdstockDecaySpec` from mapping from channels to decay functions."""
+
+    for channel in channel_function_map:
+      if channel not in self.input_data.get_all_adstock_hill_channels():
+        raise KeyError(f"Channel {channel} not found in data.")
+
+    if self.input_data.media_channel is not None:
+      media_channel_builder = self.input_data.get_paid_media_channels_argument_builder().with_default_value(
+          constants.GEOMETRIC_DECAY
+      )
+      media_adstock_function = media_channel_builder(**channel_function_map)
+    else:
+      media_adstock_function = constants.GEOMETRIC_DECAY
+
+    if self.input_data.rf_channel is not None:
+      rf_channel_builder = self.input_data.get_paid_rf_channels_argument_builder().with_default_value(
+          constants.GEOMETRIC_DECAY
+      )
+      rf_adstock_function = rf_channel_builder(**channel_function_map)
+    else:
+      rf_adstock_function = constants.GEOMETRIC_DECAY
+
+    if self.input_data.organic_media_channel is not None:
+      organic_media_channel_builder = self.input_data.get_organic_media_channels_argument_builder().with_default_value(
+          constants.GEOMETRIC_DECAY
+      )
+      organic_media_adstock_function = organic_media_channel_builder(
+          **channel_function_map
+      )
+    else:
+      organic_media_adstock_function = constants.GEOMETRIC_DECAY
+
+    if self.input_data.organic_rf_channel is not None:
+      organic_rf_channel_builder = self.input_data.get_organic_rf_channels_argument_builder().with_default_value(
+          constants.GEOMETRIC_DECAY
+      )
+      organic_rf_adstock_function = organic_rf_channel_builder(
+          **channel_function_map
+      )
+    else:
+      organic_rf_adstock_function = constants.GEOMETRIC_DECAY
+
+    return adstock_hill.AdstockDecaySpec(
+        media=media_adstock_function,
+        rf=rf_adstock_function,
+        organic_media=organic_media_adstock_function,
+        organic_rf=organic_rf_adstock_function,
+    )
+
   def _warn_setting_ignored_priors(self):
     """Raises a warning if ignored priors are set."""
     default_distribution = prior_distribution.PriorDistribution()
@@ -946,7 +1023,7 @@ class Meridian:
 
   def _check_if_no_geo_variation(
       self,
-      scaled_data: tf.Tensor,
+      scaled_data: backend.Tensor,
       data_name: str,
       data_dims: Sequence[str],
       epsilon=1e-4,
@@ -954,16 +1031,16 @@ class Meridian:
     """Raise an error if `n_knots == n_time` and data lacks geo variation."""
 
     # Result shape: [n, d], where d is the number of axes of condition.
-    col_idx_full = tf.where(tf.math.reduce_std(scaled_data, axis=0) < epsilon)[
-        :, 1
-    ]
-    col_idx_unique, _, counts = tf.unique_with_counts(col_idx_full)
+    col_idx_full = backend.get_indices_where(
+        backend.reduce_std(scaled_data, axis=0) < epsilon
+    )[:, 1]
+    col_idx_unique, _, counts = backend.unique_with_counts(col_idx_full)
     # We use the shape of scaled_data (instead of `n_time`) because the data may
     # be padded to account for lagged effects.
     data_n_time = scaled_data.shape[1]
-    mask = tf.equal(counts, data_n_time)
-    col_idx_bad = tf.boolean_mask(col_idx_unique, mask)
-    dims_bad = tf.gather(data_dims, col_idx_bad)
+    mask = backend.equal(counts, data_n_time)
+    col_idx_bad = backend.boolean_mask(col_idx_unique, mask)
+    dims_bad = backend.gather(data_dims, col_idx_bad)
 
     if col_idx_bad.shape[0] and self.knot_info.n_knots == self.n_times:
       raise ValueError(
@@ -1024,7 +1101,7 @@ class Meridian:
 
   def _check_if_no_time_variation(
       self,
-      scaled_data: tf.Tensor,
+      scaled_data: backend.Tensor,
       data_name: str,
       data_dims: Sequence[str],
       epsilon=1e-4,
@@ -1032,13 +1109,13 @@ class Meridian:
     """Raise an error if data lacks time variation."""
 
     # Result shape: [n, d], where d is the number of axes of condition.
-    col_idx_full = tf.where(tf.math.reduce_std(scaled_data, axis=1) < epsilon)[
-        :, 1
-    ]
-    col_idx_unique, _, counts = tf.unique_with_counts(col_idx_full)
-    mask = tf.equal(counts, self.n_geos)
-    col_idx_bad = tf.boolean_mask(col_idx_unique, mask)
-    dims_bad = tf.gather(data_dims, col_idx_bad)
+    col_idx_full = backend.get_indices_where(
+        backend.reduce_std(scaled_data, axis=1) < epsilon
+    )[:, 1]
+    col_idx_unique, _, counts = backend.unique_with_counts(col_idx_full)
+    mask = backend.equal(counts, self.n_geos)
+    col_idx_bad = backend.boolean_mask(col_idx_unique, mask)
+    dims_bad = backend.gather(data_dims, col_idx_bad)
     if col_idx_bad.shape[0]:
       if self.is_national:
         raise ValueError(
@@ -1058,12 +1135,19 @@ class Meridian:
             " time."
         )
 
+  def _kpi_has_variability(self):
+    """Returns True if the KPI has variability across geos and times."""
+    return self.kpi_transformer.population_scaled_stdev != 0
+
   def _validate_kpi_transformer(self):
     """Validates the KPI transformer."""
+    if self._kpi_has_variability():
+      return
+
     kpi = "kpi" if self.is_national else "population_scaled_kpi"
+
     if (
         self.n_media_channels > 0
-        and self.kpi_transformer.population_scaled_stdev == 0
         and self.model_spec.effective_media_prior_type
         in constants.PAID_MEDIA_ROI_PRIOR_TYPES
     ):
@@ -1074,7 +1158,6 @@ class Meridian:
       )
     if (
         self.n_rf_channels > 0
-        and self.kpi_transformer.population_scaled_stdev == 0
         and self.model_spec.effective_rf_prior_type
         in constants.PAID_MEDIA_ROI_PRIOR_TYPES
     ):
@@ -1082,14 +1165,44 @@ class Meridian:
           f"`{kpi}` cannot be constant with"
           f' `rf_prior_type` = "{self.model_spec.effective_rf_prior_type}".'
       )
+    if (
+        self.n_organic_media_channels > 0
+        and self.model_spec.organic_media_prior_type
+        in [constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION]
+    ):
+      raise ValueError(
+          f"`{kpi}` cannot be constant with"
+          " `organic_media_prior_type` ="
+          f' "{self.model_spec.organic_media_prior_type}".'
+      )
+    if (
+        self.n_organic_rf_channels > 0
+        and self.model_spec.organic_rf_prior_type
+        in [constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION]
+    ):
+      raise ValueError(
+          f"`{kpi}` cannot be constant with"
+          " `organic_rf_prior_type` ="
+          f' "{self.model_spec.organic_rf_prior_type}".'
+      )
+    if (
+        self.n_non_media_channels > 0
+        and self.model_spec.non_media_treatments_prior_type
+        in [constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION]
+    ):
+      raise ValueError(
+          f"`{kpi}` cannot be constant with"
+          " `non_media_treatments_prior_type` ="
+          f' "{self.model_spec.non_media_treatments_prior_type}".'
+      )
 
   def linear_predictor_counterfactual_difference_media(
       self,
-      media_transformed: tf.Tensor,
-      alpha_m: tf.Tensor,
-      ec_m: tf.Tensor,
-      slope_m: tf.Tensor,
-  ) -> tf.Tensor:
+      media_transformed: backend.Tensor,
+      alpha_m: backend.Tensor,
+      ec_m: backend.Tensor,
+      slope_m: backend.Tensor,
+  ) -> backend.Tensor:
     """Calculates linear predictor counterfactual difference for non-RF media.
 
     For non-RF media variables (paid or organic), this function calculates the
@@ -1118,18 +1231,21 @@ class Meridian:
         alpha_m,
         ec_m,
         slope_m,
+        decay_functions=self.adstock_decay_spec.media,
     )
     # Absolute values is needed because the difference is negative for mROI
     # priors and positive for ROI and contribution priors.
-    return tf.abs(media_transformed - media_transformed_counterfactual)
+    return backend.absolute(
+        media_transformed - media_transformed_counterfactual
+    )
 
   def linear_predictor_counterfactual_difference_rf(
       self,
-      rf_transformed: tf.Tensor,
-      alpha_rf: tf.Tensor,
-      ec_rf: tf.Tensor,
-      slope_rf: tf.Tensor,
-  ) -> tf.Tensor:
+      rf_transformed: backend.Tensor,
+      alpha_rf: backend.Tensor,
+      ec_rf: backend.Tensor,
+      slope_rf: backend.Tensor,
+  ) -> backend.Tensor:
     """Calculates linear predictor counterfactual difference for RF media.
 
     For RF media variables (paid or organic), this function calculates the
@@ -1159,19 +1275,20 @@ class Meridian:
         alpha=alpha_rf,
         ec=ec_rf,
         slope=slope_rf,
+        decay_functions=self.adstock_decay_spec.rf,
     )
     # Absolute values is needed because the difference is negative for mROI
     # priors and positive for ROI and contribution priors.
-    return tf.abs(rf_transformed - rf_transformed_counterfactual)
+    return backend.absolute(rf_transformed - rf_transformed_counterfactual)
 
   def calculate_beta_x(
       self,
       is_non_media: bool,
-      incremental_outcome_x: tf.Tensor,
-      linear_predictor_counterfactual_difference: tf.Tensor,
-      eta_x: tf.Tensor,
-      beta_gx_dev: tf.Tensor,
-  ) -> tf.Tensor:
+      incremental_outcome_x: backend.Tensor,
+      linear_predictor_counterfactual_difference: backend.Tensor,
+      eta_x: backend.Tensor,
+      beta_gx_dev: backend.Tensor,
+  ) -> backend.Tensor:
     """Calculates coefficient mean parameter for any treatment variable type.
 
     The "beta_x" in the function name refers to the coefficient mean parameter
@@ -1216,10 +1333,12 @@ class Meridian:
           self.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
       )
     if self.revenue_per_kpi is None:
-      revenue_per_kpi = tf.ones([self.n_geos, self.n_times], dtype=tf.float32)
+      revenue_per_kpi = backend.ones(
+          [self.n_geos, self.n_times], dtype=backend.float32
+      )
     else:
       revenue_per_kpi = self.revenue_per_kpi
-    incremental_outcome_gx_over_beta_gx = tf.einsum(
+    incremental_outcome_gx_over_beta_gx = backend.einsum(
         "...gtx,gt,g,->...gx",
         linear_predictor_counterfactual_difference,
         revenue_per_kpi,
@@ -1227,34 +1346,35 @@ class Meridian:
         self.kpi_transformer.population_scaled_stdev,
     )
     if random_effects_normal:
-      numerator_term_x = tf.einsum(
+      numerator_term_x = backend.einsum(
           "...gx,...gx,...x->...x",
           incremental_outcome_gx_over_beta_gx,
           beta_gx_dev,
           eta_x,
       )
-      denominator_term_x = tf.einsum(
+      denominator_term_x = backend.einsum(
           "...gx->...x", incremental_outcome_gx_over_beta_gx
       )
       return (incremental_outcome_x - numerator_term_x) / denominator_term_x
     # For log-normal random effects, beta_x and eta_x are not mean & std.
     # The parameterization is beta_gx ~ exp(beta_x + eta_x * N(0, 1)).
-    denominator_term_x = tf.einsum(
+    denominator_term_x = backend.einsum(
         "...gx,...gx->...x",
         incremental_outcome_gx_over_beta_gx,
-        tf.math.exp(beta_gx_dev * eta_x[..., tf.newaxis, :]),
+        backend.exp(beta_gx_dev * eta_x[..., backend.newaxis, :]),
     )
-    return tf.math.log(incremental_outcome_x) - tf.math.log(denominator_term_x)
+    return backend.log(incremental_outcome_x) - backend.log(denominator_term_x)
 
   def adstock_hill_media(
       self,
-      media: tf.Tensor,  # pylint: disable=redefined-outer-name
-      alpha: tf.Tensor,
-      ec: tf.Tensor,
-      slope: tf.Tensor,
+      media: backend.Tensor,  # pylint: disable=redefined-outer-name
+      alpha: backend.Tensor,
+      ec: backend.Tensor,
+      slope: backend.Tensor,
+      decay_functions: str | Sequence[str] = constants.GEOMETRIC_DECAY,
       n_times_output: int | None = None,
-  ) -> tf.Tensor:
-    """Transforms media using Adstock and Hill functions in the desired order.
+  ) -> backend.Tensor:
+    """Transforms media or using Adstock and Hill functions in the desired order.
 
     Args:
       media: Tensor of dimensions `(n_geos, n_media_times, n_media_channels)`
@@ -1264,6 +1384,8 @@ class Meridian:
       alpha: Uniform distribution for Adstock and Hill calculations.
       ec: Shifted half-normal distribution for Adstock and Hill calculations.
       slope: Deterministic distribution for Adstock and Hill calculations.
+      decay_functions: String or sequence of strings denoting the adstock decay
+        function(s) for each channel. Default: 'geometric'.
       n_times_output: Number of time periods to output. This argument is
         optional when the number of time periods in `media` equals
         `self.n_media_times`, in which case `n_times_output` defaults to
@@ -1284,6 +1406,7 @@ class Meridian:
         alpha=alpha,
         max_lag=self.model_spec.max_lag,
         n_times_output=n_times_output,
+        decay_functions=decay_functions,
     )
     hill_transformer = adstock_hill.HillTransformer(
         ec=ec,
@@ -1302,13 +1425,14 @@ class Meridian:
 
   def adstock_hill_rf(
       self,
-      reach: tf.Tensor,
-      frequency: tf.Tensor,
-      alpha: tf.Tensor,
-      ec: tf.Tensor,
-      slope: tf.Tensor,
+      reach: backend.Tensor,
+      frequency: backend.Tensor,
+      alpha: backend.Tensor,
+      ec: backend.Tensor,
+      slope: backend.Tensor,
+      decay_functions: str | Sequence[str] = constants.GEOMETRIC_DECAY,
       n_times_output: int | None = None,
-  ) -> tf.Tensor:
+  ) -> backend.Tensor:
     """Transforms reach and frequency (RF) using Hill and Adstock functions.
 
     Args:
@@ -1319,6 +1443,8 @@ class Meridian:
       alpha: Uniform distribution for Adstock and Hill calculations.
       ec: Shifted half-normal distribution for Adstock and Hill calculations.
       slope: Deterministic distribution for Adstock and Hill calculations.
+      decay_functions: String or sequence of strings denoting the adstock decay
+        function(s) for each channel. Default: 'geometric'.
       n_times_output: Number of time periods to output. This argument is
         optional when the number of time periods in `reach` equals
         `self.n_media_times`, in which case `n_times_output` defaults to
@@ -1343,6 +1469,7 @@ class Meridian:
         alpha=alpha,
         max_lag=self.model_spec.max_lag,
         n_times_output=n_times_output,
+        decay_functions=decay_functions,
     )
     adj_frequency = hill_transformer.forward(frequency)
     rf_out = adstock_transformer.forward(reach * adj_frequency)
@@ -1448,7 +1575,7 @@ class Meridian:
       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,