[JAX] Refactor transformers to use backend abstraction module

PiperOrigin-RevId: 795670081

Author: andyl7an

meridian/backend/__init__.py

@@ -153,6 +153,19 @@ def _jax_cast(x: Any, dtype: Any) -> "_jax.Array":
   return x.astype(dtype)
 
 
+def _jax_divide_no_nan(x, y):
+  """JAX implementation for divide_no_nan."""
+  import jax.numpy as jnp
+
+  return jnp.where(y != 0, jnp.divide(x, y), 0.0)
+
+
+def _jax_numpy_function(*args, **kwargs):  # pylint: disable=unused-argument
+  raise NotImplementedError(
+      "backend.numpy_function is not implemented for the JAX backend."
+  )
+
+
 # --- Backend Initialization ---
 _BACKEND = config.get_backend()
 
@@ -184,6 +197,7 @@ class _JaxErrors:
   stack = ops.stack
   zeros = ops.zeros
   ones = ops.ones
+  ones_like = ops.ones_like
   repeat = ops.repeat
   where = ops.where
   transpose = ops.transpose
@@ -194,6 +208,12 @@ class _JaxErrors:
   exp = ops.exp
   log = ops.log
   reduce_sum = ops.sum
+  reduce_mean = ops.mean
+  reduce_std = ops.std
+  reduce_any = ops.any
+  is_nan = ops.isnan
+  divide_no_nan = _jax_divide_no_nan
+  numpy_function = _jax_numpy_function
 
   float32 = ops.float32
   bool_ = ops.bool_
@@ -230,6 +250,7 @@ def set_random_seed(seed: int) -> None:  # pylint: disable=unused-argument
   stack = ops.stack
   zeros = ops.zeros
   ones = ops.ones
+  ones_like = ops.ones_like
   repeat = ops.repeat
   where = ops.where
   transpose = ops.transpose
@@ -240,6 +261,12 @@ def set_random_seed(seed: int) -> None:  # pylint: disable=unused-argument
   exp = ops.math.exp
   log = ops.math.log
   reduce_sum = ops.reduce_sum
+  reduce_mean = ops.reduce_mean
+  reduce_std = ops.math.reduce_std
+  reduce_any = ops.reduce_any
+  is_nan = ops.math.is_nan
+  divide_no_nan = ops.math.divide_no_nan
+  numpy_function = ops.numpy_function
 
   float32 = ops.float32
   bool_ = ops.bool

meridian/backend/test_utils.py

@@ -24,7 +24,11 @@
 
 
 def assert_allclose(
-    a: ArrayLike, b: ArrayLike, rtol: float = 1e-6, atol: float = 1e-6
+    a: ArrayLike,
+    b: ArrayLike,
+    rtol: float = 1e-6,
+    atol: float = 1e-6,
+    err_msg: str = "",
 ):
   """Backend-agnostic assertion to check if two array-like objects are close.
 
@@ -37,23 +41,55 @@ def assert_allclose(
     b: The second array-like object to compare.
     rtol: The relative tolerance parameter.
     atol: The absolute tolerance parameter.
+    err_msg: The error message to be printed in case of failure.
 
   Raises:
     AssertionError: If the two arrays are not equal within the given tolerance.
   """
-  np.testing.assert_allclose(np.array(a), np.array(b), rtol=rtol, atol=atol)
+  np.testing.assert_allclose(
+      np.array(a), np.array(b), rtol=rtol, atol=atol, err_msg=err_msg
+  )
 
 
-def assert_allequal(a: ArrayLike, b: ArrayLike):
+def assert_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
   """Backend-agnostic assertion to check if two array-like objects are equal.
 
   This function converts both inputs to NumPy arrays before comparing them.
 
   Args:
     a: The first array-like object to compare.
     b: The second array-like object to compare.
+    err_msg: The error message to be printed in case of failure.
 
   Raises:
     AssertionError: If the two arrays are not equal.
   """
-  np.testing.assert_array_equal(np.array(a), np.array(b))
+  np.testing.assert_array_equal(np.array(a), np.array(b), err_msg=err_msg)
+
+
+def assert_all_finite(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements in an array are finite.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains non-finite values.
+  """
+  if not np.all(np.isfinite(np.array(a))):
+    raise AssertionError(err_msg or "Array contains non-finite values.")
+
+
+def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements are non-negative.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains negative values.
+  """
+  if not np.all(np.array(a) >= 0):
+    raise AssertionError(err_msg or "Array contains negative values.")

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,43 +64,43 @@ def __init__(
       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 tf.reduce_any(tf.math.is_nan(self._population_scaled_median_m)):
+    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
     )
 
   @property
   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):
@@ -113,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.
 
@@ -129,25 +130,25 @@ def __init__(
         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:
@@ -161,10 +162,10 @@ def forward(
         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 (
@@ -183,8 +184,8 @@ class KpiTransformer(TensorTransformer):
 
   def __init__(
       self,
-      kpi: tf.Tensor,
-      population: tf.Tensor,
+      kpi: backend.Tensor,
+      population: backend.Tensor,
   ):
     """`KpiTransformer` constructor.
 
@@ -195,11 +196,11 @@ def __init__(
         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):
@@ -209,18 +210,18 @@ def population_scaled_mean(self):
   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]