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

meridian/model/prior_sampler.py

@@ -18,9 +18,8 @@ from collections.abc import Mapping
 from typing import TYPE_CHECKING
 
 import arviz as az
+from meridian import backend
 from meridian import constants
-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
@@ -32,8 +31,8 @@ __all__ = [
 
 
 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
@@ -51,15 +50,15 @@ 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")
 
 
 class PriorDistributionSampler:
@@ -72,7 +71,7 @@ class PriorDistributionSampler:
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the media variables.
 
     Args:
@@ -97,8 +96,8 @@ class PriorDistributionSampler:
         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),
+    beta_gm_dev = backend.tfd.Sample(
+        backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_media_channels],
         name=constants.BETA_GM_DEV,
     ).sample(**sample_kwargs)
@@ -126,6 +125,7 @@ class PriorDistributionSampler:
           alpha=media_vars[constants.ALPHA_M],
           ec=media_vars[constants.EC_M],
           slope=media_vars[constants.SLOPE_M],
+          decay_functions=mmm.adstock_decay_spec.media
       )
       linear_predictor_counterfactual_difference = (
           mmm.linear_predictor_counterfactual_difference_media(
@@ -142,20 +142,20 @@ class PriorDistributionSampler:
           eta_x=media_vars[constants.ETA_M],
           beta_gx_dev=beta_gm_dev,
       )
-      media_vars[constants.BETA_M] = tfp.distributions.Deterministic(
+      media_vars[constants.BETA_M] = backend.tfd.Deterministic(
           beta_m_value, name=constants.BETA_M
       ).sample()
 
     beta_eta_combined = (
-        media_vars[constants.BETA_M][..., tf.newaxis, :]
-        + media_vars[constants.ETA_M][..., tf.newaxis, :] * beta_gm_dev
+        media_vars[constants.BETA_M][..., backend.newaxis, :]
+        + media_vars[constants.ETA_M][..., backend.newaxis, :] * beta_gm_dev
     )
     beta_gm_value = (
         beta_eta_combined
         if mmm.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
+        else backend.exp(beta_eta_combined)
     )
-    media_vars[constants.BETA_GM] = tfp.distributions.Deterministic(
+    media_vars[constants.BETA_GM] = backend.tfd.Deterministic(
         beta_gm_value, name=constants.BETA_GM
     ).sample()
 
@@ -165,7 +165,7 @@ class PriorDistributionSampler:
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the RF variables.
 
     Args:
@@ -190,8 +190,8 @@ class PriorDistributionSampler:
         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),
+    beta_grf_dev = backend.tfd.Sample(
+        backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_rf_channels],
         name=constants.BETA_GRF_DEV,
     ).sample(**sample_kwargs)
@@ -220,6 +220,7 @@ class PriorDistributionSampler:
           alpha=rf_vars[constants.ALPHA_RF],
           ec=rf_vars[constants.EC_RF],
           slope=rf_vars[constants.SLOPE_RF],
+          decay_functions=mmm.adstock_decay_spec.rf,
       )
       linear_predictor_counterfactual_difference = (
           mmm.linear_predictor_counterfactual_difference_rf(
@@ -236,21 +237,21 @@ class PriorDistributionSampler:
           eta_x=rf_vars[constants.ETA_RF],
           beta_gx_dev=beta_grf_dev,
       )
-      rf_vars[constants.BETA_RF] = tfp.distributions.Deterministic(
+      rf_vars[constants.BETA_RF] = backend.tfd.Deterministic(
           beta_rf_value,
           name=constants.BETA_RF,
       ).sample()
 
     beta_eta_combined = (
-        rf_vars[constants.BETA_RF][..., tf.newaxis, :]
-        + rf_vars[constants.ETA_RF][..., tf.newaxis, :] * beta_grf_dev
+        rf_vars[constants.BETA_RF][..., backend.newaxis, :]
+        + rf_vars[constants.ETA_RF][..., backend.newaxis, :] * beta_grf_dev
     )
     beta_grf_value = (
         beta_eta_combined
         if mmm.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
+        else backend.exp(beta_eta_combined)
     )
-    rf_vars[constants.BETA_GRF] = tfp.distributions.Deterministic(
+    rf_vars[constants.BETA_GRF] = backend.tfd.Deterministic(
         beta_grf_value, name=constants.BETA_GRF
     ).sample()
 
@@ -260,7 +261,7 @@ class PriorDistributionSampler:
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of organic media variables.
 
     Args:
@@ -285,8 +286,8 @@ class PriorDistributionSampler:
         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),
+    beta_gom_dev = backend.tfd.Sample(
+        backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_organic_media_channels],
         name=constants.BETA_GOM_DEV,
     ).sample(**sample_kwargs)
@@ -308,6 +309,7 @@ class PriorDistributionSampler:
           alpha=organic_media_vars[constants.ALPHA_OM],
           ec=organic_media_vars[constants.EC_OM],
           slope=organic_media_vars[constants.SLOPE_OM],
+          decay_functions=mmm.adstock_decay_spec.organic_media,
       )
       beta_om_value = mmm.calculate_beta_x(
           is_non_media=False,
@@ -316,7 +318,7 @@ class PriorDistributionSampler:
           eta_x=organic_media_vars[constants.ETA_OM],
           beta_gx_dev=beta_gom_dev,
       )
-      organic_media_vars[constants.BETA_OM] = tfp.distributions.Deterministic(
+      organic_media_vars[constants.BETA_OM] = backend.tfd.Deterministic(
           beta_om_value,
           name=constants.BETA_OM,
       ).sample()
@@ -324,16 +326,16 @@ class PriorDistributionSampler:
       raise ValueError(f"Unsupported prior type: {prior_type}")
 
     beta_eta_combined = (
-        organic_media_vars[constants.BETA_OM][..., tf.newaxis, :]
-        + organic_media_vars[constants.ETA_OM][..., tf.newaxis, :]
+        organic_media_vars[constants.BETA_OM][..., backend.newaxis, :]
+        + organic_media_vars[constants.ETA_OM][..., backend.newaxis, :]
         * beta_gom_dev
     )
     beta_gom_value = (
         beta_eta_combined
         if mmm.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
+        else backend.exp(beta_eta_combined)
     )
-    organic_media_vars[constants.BETA_GOM] = tfp.distributions.Deterministic(
+    organic_media_vars[constants.BETA_GOM] = backend.tfd.Deterministic(
         beta_gom_value, name=constants.BETA_GOM
     ).sample()
 
@@ -343,7 +345,7 @@ class PriorDistributionSampler:
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Draws samples from the prior distributions of the organic RF variables.
 
     Args:
@@ -368,8 +370,8 @@ class PriorDistributionSampler:
         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),
+    beta_gorf_dev = backend.tfd.Sample(
+        backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_organic_rf_channels],
         name=constants.BETA_GORF_DEV,
     ).sample(**sample_kwargs)
@@ -392,6 +394,7 @@ class PriorDistributionSampler:
           alpha=organic_rf_vars[constants.ALPHA_ORF],
           ec=organic_rf_vars[constants.EC_ORF],
           slope=organic_rf_vars[constants.SLOPE_ORF],
+          decay_functions=mmm.adstock_decay_spec.organic_rf,
       )
       beta_orf_value = mmm.calculate_beta_x(
           is_non_media=False,
@@ -400,7 +403,7 @@ class PriorDistributionSampler:
           eta_x=organic_rf_vars[constants.ETA_ORF],
           beta_gx_dev=beta_gorf_dev,
       )
-      organic_rf_vars[constants.BETA_ORF] = tfp.distributions.Deterministic(
+      organic_rf_vars[constants.BETA_ORF] = backend.tfd.Deterministic(
           beta_orf_value,
           name=constants.BETA_ORF,
       ).sample()
@@ -408,15 +411,16 @@ class PriorDistributionSampler:
       raise ValueError(f"Unsupported prior type: {prior_type}")
 
     beta_eta_combined = (
-        organic_rf_vars[constants.BETA_ORF][..., tf.newaxis, :]
-        + organic_rf_vars[constants.ETA_ORF][..., tf.newaxis, :] * beta_gorf_dev
+        organic_rf_vars[constants.BETA_ORF][..., backend.newaxis, :]
+        + organic_rf_vars[constants.ETA_ORF][..., backend.newaxis, :]
+        * beta_gorf_dev
     )
     beta_gorf_value = (
         beta_eta_combined
         if mmm.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
-        else tf.math.exp(beta_eta_combined)
+        else backend.exp(beta_eta_combined)
     )
-    organic_rf_vars[constants.BETA_GORF] = tfp.distributions.Deterministic(
+    organic_rf_vars[constants.BETA_GORF] = backend.tfd.Deterministic(
         beta_gorf_value, name=constants.BETA_GORF
     ).sample()
 
@@ -426,7 +430,7 @@ class PriorDistributionSampler:
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Draws from the prior distributions of the non-media treatment variables.
 
     Args:
@@ -448,8 +452,8 @@ class PriorDistributionSampler:
     non_media_treatments_vars = {
         constants.XI_N: prior.xi_n.sample(**sample_kwargs),
     }
-    gamma_gn_dev = tfp.distributions.Sample(
-        tfp.distributions.Normal(0, 1),
+    gamma_gn_dev = backend.tfd.Sample(
+        backend.tfd.Normal(0, 1),
         [mmm.n_geos, mmm.n_non_media_channels],
         name=constants.GAMMA_GN_DEV,
     ).sample(**sample_kwargs)
@@ -479,35 +483,31 @@ class PriorDistributionSampler:
           eta_x=non_media_treatments_vars[constants.XI_N],
           beta_gx_dev=gamma_gn_dev,
       )
-      non_media_treatments_vars[constants.GAMMA_N] = (
-          tfp.distributions.Deterministic(
-              gamma_n_value, name=constants.GAMMA_N
-          ).sample()
-      )
+      non_media_treatments_vars[constants.GAMMA_N] = backend.tfd.Deterministic(
+          gamma_n_value, name=constants.GAMMA_N
+      ).sample()
     else:
       raise ValueError(f"Unsupported prior type: {prior_type}")
-    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()
-    )
+    non_media_treatments_vars[constants.GAMMA_GN] = backend.tfd.Deterministic(
+        non_media_treatments_vars[constants.GAMMA_N][..., backend.newaxis, :]
+        + non_media_treatments_vars[constants.XI_N][..., backend.newaxis, :]
+        * gamma_gn_dev,
+        name=constants.GAMMA_GN,
+    ).sample()
     return non_media_treatments_vars
 
   def _sample_prior(
       self,
       n_draws: int,
       seed: int | None = None,
-  ) -> Mapping[str, tf.Tensor]:
+  ) -> Mapping[str, backend.Tensor]:
     """Returns a mapping of prior parameters to tensors of the samples."""
     mmm = self._meridian
 
     # 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)
+      backend.set_random_seed(seed)
 
     prior = mmm.prior_broadcast
     # `sample_shape` is prepended to the shape of each BatchBroadcast in `prior`
@@ -527,11 +527,11 @@ class PriorDistributionSampler:
         ),
     }
 
-    base_vars[constants.MU_T] = tfp.distributions.Deterministic(
-        tf.einsum(
+    base_vars[constants.MU_T] = backend.tfd.Deterministic(
+        backend.einsum(
             "...k,kt->...t",
             base_vars[constants.KNOT_VALUES],
-            tf.convert_to_tensor(mmm.knot_info.weights),
+            backend.to_tensor(mmm.knot_info.weights),
         ),
         name=constants.MU_T,
     ).sample()
@@ -544,14 +544,14 @@ class PriorDistributionSampler:
           constants.XI_C: prior.xi_c.sample(**sample_kwargs),
       }
 
-      gamma_gc_dev = tfp.distributions.Sample(
-          tfp.distributions.Normal(0, 1),
+      gamma_gc_dev = backend.tfd.Sample(
+          backend.tfd.Normal(0, 1),
           [mmm.n_geos, mmm.n_controls],
           name=constants.GAMMA_GC_DEV,
       ).sample(**sample_kwargs)
-      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,
+      base_vars[constants.GAMMA_GC] = backend.tfd.Deterministic(
+          base_vars[constants.GAMMA_C][..., backend.newaxis, :]
+          + base_vars[constants.XI_C][..., backend.newaxis, :] * gamma_gc_dev,
           name=constants.GAMMA_GC,
       ).sample()
 

meridian/model/spec.py

@@ -14,6 +14,7 @@
 
 """Defines model specification parameters for Meridian."""
 
+from collections.abc import Mapping
 import dataclasses
 from typing import Sequence
 import warnings
@@ -72,8 +73,7 @@ class ModelSpec:
       before Hill. This argument does not apply to RF channels. Default:
       `False`.
     max_lag: An integer indicating the maximum number of lag periods (≥ `0`) to
-      include in the Adstock calculation. Can also be set to `None`, which is
-      equivalent to infinite max lag. Default: `8`.
+      include in the Adstock calculation. Default: `8`.
     unique_sigma_for_each_geo: A boolean indicating whether to use a unique
       residual variance for each geo. If `False`, then a single residual
       variance is used for all geos. Default: `False`.
@@ -202,6 +202,20 @@ class ModelSpec:
       `(n_non_media_channels,)` indicating the non-media variables for which the
       non-media value will be scaled by population. If `None`, then no non-media
       variables are scaled by population. Default: `None`.
+    adstock_decay_spec: A string or mapping specifying the adstock decay
+      function for each media, RF, organic media and organic RF channel.
+      * If a string, must be either `'geometric'` or `'binomial'`, specifying
+        that decay function for all channels.
+      * If a mapping, keys should be channel names and values should be
+        `'geometric'` or `'binomial'`, with each key-value pair denoting the
+        adstock decay function to use for that channel. Channels that are not
+        specified in the mapping default to using 'geometric'.
+      Default: `'geometric'`.
+    enable_aks: A boolean indicating whether to use the Automatic Knot Selection
+      algorithm to select an optimal number of knots for running the model
+      instead of the default 1 for national models and n_times for geo models.
+      If this is set to `True` and the `knots` arg is provided, then an error
+      will be raised. Default: `False`.
   """
 
   prior: prior_distribution.PriorDistribution = dataclasses.field(
@@ -209,7 +223,7 @@ class ModelSpec:
   )
   media_effects_dist: str = constants.MEDIA_EFFECTS_LOG_NORMAL
   hill_before_adstock: bool = False
-  max_lag: int | None = 8
+  max_lag: int = 8
   unique_sigma_for_each_geo: bool = False
   media_prior_type: str | None = None
   rf_prior_type: str | None = None
@@ -227,6 +241,8 @@ class ModelSpec:
   holdout_id: np.ndarray | None = None
   control_population_scaling_id: np.ndarray | None = None
   non_media_population_scaling_id: np.ndarray | None = None
+  adstock_decay_spec: str | Mapping[str, str] = constants.GEOMETRIC_DECAY
+  enable_aks: bool = False
 
   def __post_init__(self):
     # Validate media_effects_dist.
@@ -311,6 +327,10 @@ class ModelSpec:
       raise ValueError("The `knots` parameter cannot be an empty list.")
     if isinstance(self.knots, int) and self.knots == 0:
       raise ValueError("The `knots` parameter cannot be zero.")
+    if self.knots is not None and self.enable_aks:
+      raise ValueError(
+          "The `knots` parameter cannot be set when `enable_aks` is True."
+      )
 
   @property
   def effective_media_prior_type(self) -> str:

meridian/model/transformers.py

@@ -15,8 +15,9 @@
 """Contains data transformers for various inputs of the Meridian model."""
 
 import abc
+
+from meridian import backend
 import numpy as np
-import tensorflow as tf
 
 
 __all__ = [
@@ -31,14 +32,14 @@ class TensorTransformer(abc.ABC):
   """Abstract class for data transformers."""
 
   @abc.abstractmethod
-  @tf.function(jit_compile=True)
-  def forward(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def forward(self, tensor: backend.Tensor) -> backend.Tensor:
     """Transforms a given tensor."""
     raise NotImplementedError("`forward` must be implemented.")
 
   @abc.abstractmethod
-  @tf.function(jit_compile=True)
-  def inverse(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def inverse(self, tensor: backend.Tensor) -> backend.Tensor:
     """Transforms back a given tensor."""
     raise NotImplementedError("`inverse` must be implemented.")
 
@@ -52,8 +53,8 @@ class MediaTransformer(TensorTransformer):
 
   def __init__(
       self,
-      media: tf.Tensor,
-      population: tf.Tensor,
+      media: backend.Tensor,
+      population: backend.Tensor,
   ):
     """`MediaTransformer` constructor.
 
@@ -63,22 +64,27 @@ class MediaTransformer(TensorTransformer):
       population: A tensor of dimension `(n_geos,)` containing the population of
         each geo, used to compute the scale factors.
     """
-    population_scaled_media = tf.math.divide_no_nan(
-        media, population[:, tf.newaxis, tf.newaxis]
+    population_scaled_media = backend.divide_no_nan(
+        media, population[:, backend.newaxis, backend.newaxis]
     )
     # Replace zeros with NaNs
-    population_scaled_media_nan = tf.where(
+    population_scaled_media_nan = backend.where(
         population_scaled_media == 0, np.nan, population_scaled_media
     )
     # Tensor of medians of the positive portion of `media`. Used as a component
     # for scaling.
-    self._population_scaled_median_m = tf.numpy_function(
+    self._population_scaled_median_m = backend.numpy_function(
         func=lambda x: np.nanmedian(x, axis=[0, 1]),
         inp=[population_scaled_media_nan],
-        Tout=tf.float32,
+        Tout=backend.float32,
     )
+    if backend.reduce_any(backend.is_nan(self._population_scaled_median_m)):
+      raise ValueError(
+          "MediaTransformer has a NaN population-scaled non-zero median due to"
+          " a media channel with either all zeroes or all NaNs."
+      )
     # Tensor of dimensions (`n_geos` x 1) of weights for scaling `metric`.
-    self._scale_factors_gm = tf.einsum(
+    self._scale_factors_gm = backend.einsum(
         "g,m->gm", population, self._population_scaled_median_m
     )
 
@@ -86,15 +92,15 @@ class MediaTransformer(TensorTransformer):
   def population_scaled_median_m(self):
     return self._population_scaled_median_m
 
-  @tf.function(jit_compile=True)
-  def forward(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def forward(self, tensor: backend.Tensor) -> backend.Tensor:
     """Scales a given tensor using the stored scale factors."""
-    return tensor / self._scale_factors_gm[:, tf.newaxis, :]
+    return tensor / self._scale_factors_gm[:, backend.newaxis, :]
 
-  @tf.function(jit_compile=True)
-  def inverse(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def inverse(self, tensor: backend.Tensor) -> backend.Tensor:
     """Scales a given tensor using the inversed stored scale factors."""
-    return tensor * self._scale_factors_gm[:, tf.newaxis, :]
+    return tensor * self._scale_factors_gm[:, backend.newaxis, :]
 
 
 class CenteringAndScalingTransformer(TensorTransformer):
@@ -108,9 +114,9 @@ class CenteringAndScalingTransformer(TensorTransformer):
 
   def __init__(
       self,
-      tensor: tf.Tensor,
-      population: tf.Tensor,
-      population_scaling_id: tf.Tensor | None = None,
+      tensor: backend.Tensor,
+      population: backend.Tensor,
+      population_scaling_id: backend.Tensor | None = None,
   ):
     """`CenteringAndScalingTransformer` constructor.
 
@@ -124,25 +130,25 @@ class CenteringAndScalingTransformer(TensorTransformer):
         scaled by population.
     """
     if population_scaling_id is not None:
-      self._population_scaling_factors = tf.where(
+      self._population_scaling_factors = backend.where(
           population_scaling_id,
           population[:, None],
-          tf.ones_like(population)[:, None],
+          backend.ones_like(population)[:, None],
       )
       population_scaled_tensor = (
           tensor / self._population_scaling_factors[:, None, :]
       )
-      self._means = tf.reduce_mean(population_scaled_tensor, axis=(0, 1))
-      self._stdevs = tf.math.reduce_std(population_scaled_tensor, axis=(0, 1))
+      self._means = backend.reduce_mean(population_scaled_tensor, axis=(0, 1))
+      self._stdevs = backend.reduce_std(population_scaled_tensor, axis=(0, 1))
     else:
       self._population_scaling_factors = None
-      self._means = tf.reduce_mean(tensor, axis=(0, 1))
-      self._stdevs = tf.math.reduce_std(tensor, axis=(0, 1))
+      self._means = backend.reduce_mean(tensor, axis=(0, 1))
+      self._stdevs = backend.reduce_std(tensor, axis=(0, 1))
 
-  @tf.function(jit_compile=True)
+  @backend.function(jit_compile=True)
   def forward(
-      self, tensor: tf.Tensor, apply_population_scaling: bool = True
-  ) -> tf.Tensor:
+      self, tensor: backend.Tensor, apply_population_scaling: bool = True
+  ) -> backend.Tensor:
     """Scales a given tensor using the stored coefficients.
 
     Args:
@@ -156,10 +162,10 @@ class CenteringAndScalingTransformer(TensorTransformer):
         and self._population_scaling_factors is not None
     ):
       tensor /= self._population_scaling_factors[:, None, :]
-    return tf.math.divide_no_nan(tensor - self._means, self._stdevs)
+    return backend.divide_no_nan(tensor - self._means, self._stdevs)
 
-  @tf.function(jit_compile=True)
-  def inverse(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def inverse(self, tensor: backend.Tensor) -> backend.Tensor:
     """Scales back a given tensor using the stored coefficients."""
     scaled_tensor = tensor * self._stdevs + self._means
     return (
@@ -178,8 +184,8 @@ class KpiTransformer(TensorTransformer):
 
   def __init__(
       self,
-      kpi: tf.Tensor,
-      population: tf.Tensor,
+      kpi: backend.Tensor,
+      population: backend.Tensor,
   ):
     """`KpiTransformer` constructor.
 
@@ -190,11 +196,11 @@ class KpiTransformer(TensorTransformer):
         each geo, used to to compute the population scale factors.
     """
     self._population = population
-    population_scaled_kpi = tf.math.divide_no_nan(
-        kpi, self._population[:, tf.newaxis]
+    population_scaled_kpi = backend.divide_no_nan(
+        kpi, self._population[:, backend.newaxis]
     )
-    self._population_scaled_mean = tf.reduce_mean(population_scaled_kpi)
-    self._population_scaled_stdev = tf.math.reduce_std(population_scaled_kpi)
+    self._population_scaled_mean = backend.reduce_mean(population_scaled_kpi)
+    self._population_scaled_stdev = backend.reduce_std(population_scaled_kpi)
 
   @property
   def population_scaled_mean(self):
@@ -204,18 +210,18 @@ class KpiTransformer(TensorTransformer):
   def population_scaled_stdev(self):
     return self._population_scaled_stdev
 
-  @tf.function(jit_compile=True)
-  def forward(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def forward(self, tensor: backend.Tensor) -> backend.Tensor:
     """Scales a given tensor using the stored coefficients."""
-    return tf.math.divide_no_nan(
-        tf.math.divide_no_nan(tensor, self._population[:, tf.newaxis])
+    return backend.divide_no_nan(
+        backend.divide_no_nan(tensor, self._population[:, backend.newaxis])
         - self._population_scaled_mean,
         self._population_scaled_stdev,
     )
 
-  @tf.function(jit_compile=True)
-  def inverse(self, tensor: tf.Tensor) -> tf.Tensor:
+  @backend.function(jit_compile=True)
+  def inverse(self, tensor: backend.Tensor) -> backend.Tensor:
     """Scales back a given tensor using the stored coefficients."""
     return (
         tensor * self._population_scaled_stdev + self._population_scaled_mean
-    ) * self._population[:, tf.newaxis]
+    ) * self._population[:, backend.newaxis]

meridian/version.py

@@ -14,4 +14,4 @@
 
 """Module for Meridian version."""
 
-__version__ = "1.1.5"
+__version__ = "1.2.0"