google-meridian 1.3.0__py3-none-any.whl → 1.3.2__py3-none-any.whl

meridian/backend/__init__.py

@@ -19,6 +19,7 @@ import functools
 import os
 from typing import Any, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 import warnings
+
 from meridian.backend import config
 import numpy as np
 from typing_extensions import Literal
@@ -220,7 +221,7 @@ def _tf_arange(
 
 def _jax_cast(x: Any, dtype: Any) -> "_jax.Array":
   """JAX implementation for cast."""
-  return x.astype(dtype)
+  return jax_ops.asarray(x, dtype=dtype)
 
 
 def _jax_divide_no_nan(x, y):
@@ -305,17 +306,132 @@ def _jax_numpy_function(*args, **kwargs):  # pylint: disable=unused-argument
   )
 
 
-def _jax_make_tensor_proto(*args, **kwargs):  # pylint: disable=unused-argument
-  raise NotImplementedError(
-      "backend.make_tensor_proto is not implemented for the JAX backend."
-  )
+def _jax_make_tensor_proto(values, dtype=None, shape=None):  # pylint: disable=unused-argument
+  """JAX implementation for make_tensor_proto."""
+  # pylint: disable=g-direct-tensorflow-import
+  from tensorflow.core.framework import tensor_pb2
+  from tensorflow.core.framework import tensor_shape_pb2
+  from tensorflow.core.framework import types_pb2
+  # pylint: enable=g-direct-tensorflow-import
 
+  if not isinstance(values, np.ndarray):
+    values = np.array(values)
 
-def _jax_make_ndarray(*args, **kwargs):  # pylint: disable=unused-argument
-  raise NotImplementedError(
-      "backend.make_ndarray is not implemented for the JAX backend."
+  if dtype:
+    numpy_dtype = np.dtype(dtype)
+    values = values.astype(numpy_dtype)
+  else:
+    numpy_dtype = values.dtype
+
+  dtype_map = {
+      np.dtype(np.float16): types_pb2.DT_HALF,
+      np.dtype(np.float32): types_pb2.DT_FLOAT,
+      np.dtype(np.float64): types_pb2.DT_DOUBLE,
+      np.dtype(np.int32): types_pb2.DT_INT32,
+      np.dtype(np.uint8): types_pb2.DT_UINT8,
+      np.dtype(np.uint16): types_pb2.DT_UINT16,
+      np.dtype(np.uint32): types_pb2.DT_UINT32,
+      np.dtype(np.uint64): types_pb2.DT_UINT64,
+      np.dtype(np.int16): types_pb2.DT_INT16,
+      np.dtype(np.int8): types_pb2.DT_INT8,
+      np.dtype(np.int64): types_pb2.DT_INT64,
+      np.dtype(np.complex64): types_pb2.DT_COMPLEX64,
+      np.dtype(np.complex128): types_pb2.DT_COMPLEX128,
+      np.dtype(np.bool_): types_pb2.DT_BOOL,
+      # Note: String types are handled outside the map.
+  }
+  proto_dtype = dtype_map.get(numpy_dtype)
+  if proto_dtype is None and numpy_dtype.kind in ("S", "U"):
+    proto_dtype = types_pb2.DT_STRING
+
+  if proto_dtype is None:
+    raise TypeError(
+        f"Unsupported dtype for TensorProto conversion: {numpy_dtype}"
+    )
+
+  proto = tensor_pb2.TensorProto(
+      dtype=proto_dtype,
+      tensor_shape=tensor_shape_pb2.TensorShapeProto(
+          dim=[
+              tensor_shape_pb2.TensorShapeProto.Dim(size=d)
+              for d in values.shape
+          ]
+      ),
   )
 
+  proto.tensor_content = values.tobytes()
+  return proto
+
+
+def _jax_make_ndarray(proto):
+  """JAX implementation for make_ndarray."""
+  # pylint: disable=g-direct-tensorflow-import
+  from tensorflow.core.framework import types_pb2
+  # pylint: enable=g-direct-tensorflow-import
+
+  dtype_map = {
+      types_pb2.DT_HALF: np.float16,
+      types_pb2.DT_FLOAT: np.float32,
+      types_pb2.DT_DOUBLE: np.float64,
+      types_pb2.DT_INT32: np.int32,
+      types_pb2.DT_UINT8: np.uint8,
+      types_pb2.DT_UINT16: np.uint16,
+      types_pb2.DT_UINT32: np.uint32,
+      types_pb2.DT_UINT64: np.uint64,
+      types_pb2.DT_INT16: np.int16,
+      types_pb2.DT_INT8: np.int8,
+      types_pb2.DT_INT64: np.int64,
+      types_pb2.DT_COMPLEX64: np.complex64,
+      types_pb2.DT_COMPLEX128: np.complex128,
+      types_pb2.DT_BOOL: np.bool_,
+      types_pb2.DT_STRING: np.bytes_,
+  }
+  if proto.dtype not in dtype_map:
+    raise TypeError(f"Unsupported TensorProto dtype: {proto.dtype}")
+
+  shape = [d.size for d in proto.tensor_shape.dim]
+  dtype = dtype_map[proto.dtype]
+
+  if proto.tensor_content:
+    num_elements = np.prod(shape).item() if shape else 0
+    # When deserializing a string from tensor_content, the itemsize is not
+    # explicitly stored. We must infer it from the content length and shape.
+    if dtype == np.bytes_ and num_elements > 0:
+      content_len = len(proto.tensor_content)
+      itemsize = content_len // num_elements
+      if itemsize * num_elements != content_len:
+        raise ValueError(
+            "Tensor content size is not a multiple of the number of elements"
+            " for string dtype."
+        )
+      dtype = np.dtype(f"S{itemsize}")
+
+    return (
+        np.frombuffer(proto.tensor_content, dtype=dtype).copy().reshape(shape)
+    )
+
+  # Fallback for protos that store data in val fields instead of tensor_content.
+  if dtype == np.float32:
+    val_field = proto.float_val
+  elif dtype == np.float64:
+    val_field = proto.double_val
+  elif dtype == np.int32:
+    val_field = proto.int_val
+  elif dtype == np.int64:
+    val_field = proto.int64_val
+  elif dtype == np.bool_:
+    val_field = proto.bool_val
+  else:
+    if proto.string_val:
+      return np.array(proto.string_val, dtype=np.bytes_).reshape(shape)
+    if not any(shape):
+      return np.array([], dtype=dtype).reshape(shape)
+    raise TypeError(
+        f"Unsupported dtype for TensorProto value field fallback: {dtype}"
+    )
+
+  return np.array(val_field, dtype=dtype).reshape(shape)
+
 
 def _jax_get_indices_where(condition):
   """JAX implementation for get_indices_where."""
@@ -493,16 +609,75 @@ def _tf_get_seed_data(seed: Any) -> Optional[np.ndarray]:
 
 
 def _jax_convert_to_tensor(data, dtype=None):
-  """Converts data to a JAX array, handling strings as NumPy arrays."""
+  """Converts data to a JAX array, handling strings as NumPy arrays.
+
+  This function explicitly unwraps objects with a `.values` attribute (e.g.,
+  pandas.DataFrame, xarray.DataArray) to access the underlying NumPy array,
+  provided that `.values` is not a method. This takes precedence over the
+  `__array__` protocol.
+
+  It also handles precision mismatches: if `data` is float64 and `dtype` is
+  not specified, and JAX x64 mode is disabled (default), it issues a warning
+  and explicitly casts to float32 to match the backend default and prevent
+  silent precision loss or type errors in downstream operations.
+
+  Args:
+    data: The data to convert.
+    dtype: The desired data type.
+
+  Returns:
+    A JAX array, or a NumPy array if the dtype is a string type.
+  """
+  # Unwrap xarray.DataArray, pandas.Series, and pandas.DataFrame objects.
+  # These objects wrap the underlying NumPy array in a .values attribute.
+  if hasattr(data, "values") and not callable(data.values):
+    data = data.values
+
+  # Convert to numpy array upfront to simplify dtype inspection below.
+  # A standard Python float is 64-bit, and this conversion allows the
+  # subsequent logic to correctly detect and handle potential float64
+  # downcasting for scalar inputs.
+  if isinstance(data, (list, tuple, float)):
+    data = np.array(data)
+
   # JAX does not natively support string tensors in the same way TF does.
   # If a string dtype is requested, or if the data is inherently strings,
   # we fall back to a standard NumPy array.
-  if dtype == np.str_ or (
-      dtype is None
-      and isinstance(data, (list, np.ndarray))
-      and np.array(data).dtype.kind in ("S", "U")
-  ):
-    return np.array(data, dtype=np.str_)
+  is_string_target = False
+  if dtype is not None:
+    try:
+      if np.dtype(dtype).kind in ("S", "U"):
+        is_string_target = True
+    except TypeError:
+      # This can happen if dtype is not a valid dtype specifier,
+      # let jax.asarray handle it.
+      pass
+
+  is_string_data = isinstance(data, np.ndarray) and data.dtype.kind in (
+      "S",
+      "U",
+  )
+
+  if is_string_target:
+    return np.array(data, dtype=dtype)
+
+  if dtype is None and is_string_data:
+    return data
+
+  # If the user provides float64 data but does not request a specific dtype,
+  # and JAX 64-bit mode is disabled (default), JAX would implicitly truncate.
+  # We cast to float32 and warn the user to prevent silent mismatches.
+  if dtype is None:
+    is_float64_input = hasattr(data, "dtype") and data.dtype == np.float64
+    if is_float64_input:
+      if not jax.config.jax_enable_x64:
+        warnings.warn(
+            "Input data is float64. Casting to float32 to match backend "
+            "default precision.",
+            UserWarning,
+        )
+        dtype = jax_ops.float32
+
   return jax_ops.asarray(data, dtype=dtype)
 
 
@@ -535,18 +710,48 @@ def _tf_nanvar(a, axis=None, keepdims=False):
   return tf.convert_to_tensor(var)
 
 
-def _jax_one_hot(*args, **kwargs):  # pylint: disable=unused-argument
+def _jax_one_hot(
+    indices, depth, on_value=None, off_value=None, axis=None, dtype=None
+):
   """JAX implementation for one_hot."""
-  raise NotImplementedError(
-      "backend.one_hot is not implemented for the JAX backend."
+  import jax.numpy as jnp
+
+  resolved_dtype = _resolve_dtype(dtype, on_value, off_value, 1, 0)
+  jax_axis = -1 if axis is None else axis
+
+  one_hot_result = jax.nn.one_hot(
+      indices, num_classes=depth, dtype=jnp.dtype(resolved_dtype), axis=jax_axis
   )
 
+  on_val = 1 if on_value is None else on_value
+  off_val = 0 if off_value is None else off_value
+
+  if on_val == 1 and off_val == 0:
+    return one_hot_result
 
-def _jax_roll(*args, **kwargs):  # pylint: disable=unused-argument
+  on_tensor = jnp.array(on_val, dtype=jnp.dtype(resolved_dtype))
+  off_tensor = jnp.array(off_val, dtype=jnp.dtype(resolved_dtype))
+
+  return jnp.where(one_hot_result == 1, on_tensor, off_tensor)
+
+
+def _jax_roll(a, shift, axis=None):
   """JAX implementation for roll."""
-  raise NotImplementedError(
-      "backend.roll is not implemented for the JAX backend."
-  )
+  import jax.numpy as jnp
+
+  return jnp.roll(a, shift, axis=axis)
+
+
+def _tf_roll(a, shift: Sequence[int], axis=None):
+  """TensorFlow implementation for roll that handles axis=None."""
+  import tensorflow as tf
+
+  if axis is None:
+    original_shape = tf.shape(a)
+    flat_tensor = tf.reshape(a, [-1])
+    rolled_flat = tf.roll(flat_tensor, shift=shift, axis=0)
+    return tf.reshape(rolled_flat, original_shape)
+  return tf.roll(a, shift, axis=axis)
 
 
 def _jax_enable_op_determinism():
@@ -772,7 +977,7 @@ if _BACKEND == config.Backend.JAX:
   newaxis = _ops.newaxis
   TensorShape = _jax_tensor_shape
   int32 = _ops.int32
-  string = np.str_
+  string = np.bytes_
 
   stabilize_rf_roi_grid = _jax_stabilize_rf_roi_grid
 
@@ -904,7 +1109,7 @@ elif _BACKEND == config.Backend.TENSORFLOW:
   reduce_sum = _ops.reduce_sum
   repeat = _ops.repeat
   reshape = _ops.reshape
-  roll = _ops.roll
+  roll = _tf_roll
   set_random_seed = tf_backend.keras.utils.set_random_seed
   split = _ops.split
   stack = _ops.stack

meridian/backend/test_utils.py

@@ -14,12 +14,25 @@
 
 """Common testing utilities for Meridian, designed to be backend-agnostic."""
 
+import dataclasses
 from typing import Any, Optional
+
 from absl.testing import parameterized
+from google.protobuf import descriptor
+from google.protobuf import message
 from meridian import backend
 from meridian.backend import config
 import numpy as np
 
+from tensorflow.python.util.protobuf import compare
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.core.framework import tensor_pb2
+
+
+# pylint: enable=g-direct-tensorflow-import
+
+FieldDescriptor = descriptor.FieldDescriptor
+
 # A type alias for backend-agnostic array-like objects.
 # We use `Any` here to avoid circular dependencies with the backend module
 # while still allowing the function to accept backend-specific tensor types.
@@ -70,6 +83,75 @@ def assert_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
   np.testing.assert_array_equal(np.array(a), np.array(b), err_msg=err_msg)
 
 
+def assert_deep_equals(
+    test_case,
+    obj1: Any,
+    obj2: Any,
+    msg: str = "",
+    rtol: float = 1e-5,
+    atol: float = 1e-5,
+):
+  """Recursive equality check handling Dataclasses, Lists, and Backend Tensors.
+
+  Args:
+    test_case: The unittest.TestCase instance (self) to use for assertions.
+    obj1: The first object to compare.
+    obj2: The second object to compare.
+    msg: Optional error message prefix.
+    rtol: Relative tolerance for float comparison.
+    atol: Absolute tolerance for float comparison.
+  """
+  if obj1 is None or obj2 is None:
+    test_case.assertEqual(obj1, obj2, msg=msg)
+    return
+
+  if (
+      hasattr(obj1, "__array__")
+      or hasattr(obj1, "numpy")
+      or isinstance(obj1, (np.ndarray, backend.Tensor))
+  ):
+    arr1 = np.array(obj1)
+    arr2 = np.array(obj2)
+
+    # Check for non-numeric types where atol/rtol don't apply
+    if arr1.dtype.kind in ("U", "S", "O", "b"):
+      np.testing.assert_array_equal(arr1, arr2, err_msg=msg)
+    else:
+      np.testing.assert_allclose(arr1, arr2, err_msg=msg, rtol=rtol, atol=atol)
+    return
+
+  if dataclasses.is_dataclass(obj1):
+    test_case.assertIs(
+        type(obj1),
+        type(obj2),
+        msg=f"{msg} Type mismatch: {type(obj1)} vs {type(obj2)}",
+    )
+    for field in dataclasses.fields(obj1):
+      val1 = getattr(obj1, field.name)
+      val2 = getattr(obj2, field.name)
+      assert_deep_equals(
+          test_case,
+          val1,
+          val2,
+          msg=f"{msg}.{field.name}",
+          rtol=rtol,
+          atol=atol,
+      )
+    return
+
+  if isinstance(obj1, (list, tuple)):
+    test_case.assertIsInstance(obj2, (list, tuple), msg=f"{msg} Type mismatch")
+    test_case.assertEqual(len(obj1), len(obj2), msg=f"{msg} Length mismatch")
+    for i, (item1, item2) in enumerate(zip(obj1, obj2)):
+      assert_deep_equals(
+          test_case, item1, item2, msg=f"{msg}[{i}]", rtol=rtol, atol=atol
+      )
+    return
+
+  # Fallback to standard equality for primitives (int, str, float, etc.)
+  test_case.assertEqual(obj1, obj2, msg=msg)
+
+
 def assert_seed_allequal(a: Any, b: Any, err_msg: str = ""):
   """Backend-agnostic assertion to check if two seed objects are equal."""
   data_a = backend.get_seed_data(a)
@@ -131,6 +213,118 @@ def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
     raise AssertionError(err_msg or "Array contains negative values.")
 
 
+# --- Proto Utilities ---
+def normalize_tensor_protos(proto: message.Message):
+  """Recursively normalizes TensorProto messages within a proto (In-place).
+
+  This ensures a consistent serialization format across different backends
+  (e.g., JAX vs TF) by repacking TensorProtos using the current backend's
+  canonical method (backend.make_tensor_proto). This handles differences
+  like using `bool_val` versus `tensor_content` for boolean tensors.
+
+  Args:
+    proto: The protobuf message object to normalize. This object is modified in
+      place.
+  """
+  if not isinstance(proto, message.Message):
+    return
+
+  for desc, value in proto.ListFields():
+    if desc.type != FieldDescriptor.TYPE_MESSAGE:
+      continue
+
+    # A map is defined as a repeated field whose message type has the
+    # map_entry option set.
+    is_map = (
+        desc.label == FieldDescriptor.LABEL_REPEATED
+        and desc.message_type.has_options
+        and desc.message_type.GetOptions().map_entry
+    )
+
+    if is_map:
+      for item in value.values():
+        # Helper checks if values are scalars or messages.
+        _process_message_for_normalization(item)
+
+    elif desc.label == FieldDescriptor.LABEL_REPEATED:
+      # Handle standard repeated message fields.
+      for item in value:
+        _process_message_for_normalization(item)
+    else:
+      # Handle singular message fields.
+      _process_message_for_normalization(value)
+
+
+def _process_message_for_normalization(msg: Any):
+  """Helper to process a potential message during normalization traversal."""
+  # Ensure we only process message objects.
+  # If msg is a scalar (e.g., string from map<string, string>), stop recursion.
+  if not isinstance(msg, message.Message):
+    return
+
+  if isinstance(msg, tensor_pb2.TensorProto):
+    _repack_tensor_proto(msg)
+  else:
+    # If it's another message type, recurse into its fields.
+    normalize_tensor_protos(msg)
+
+
+def _repack_tensor_proto(tensor_proto: "tensor_pb2.TensorProto"):
+  """Repacks a TensorProto in place to use a consistent serialization format."""
+  if not tensor_proto.ByteSize():
+    return
+
+  try:
+    data_array = backend.make_ndarray(tensor_proto)
+  except Exception as e:
+    raise ValueError(
+        "Failed to deserialize TensorProto during normalization:"
+        f" {e}\nProto content:\n{tensor_proto}"
+    ) from e
+
+  new_tensor_proto = backend.make_tensor_proto(data_array)
+
+  tensor_proto.Clear()
+  tensor_proto.CopyFrom(new_tensor_proto)
+
+
+def assert_normalized_proto_equal(
+    test_case: parameterized.TestCase,
+    expected: message.Message,
+    actual: message.Message,
+    msg: Optional[str] = None,
+    **kwargs: Any,
+):
+  """Compares two protos after normalizing TensorProto fields.
+
+  Use this instead of compare.assertProtoEqual when protos contain tensors
+  to ensure backend-agnostic comparison.
+
+  Args:
+    test_case: The TestCase instance (self).
+    expected: The expected protobuf message.
+    actual: The actual protobuf message.
+    msg: An optional message to display on failure.
+    **kwargs: Additional keyword arguments passed to assertProto2Equal (e.g.,
+      precision).
+  """
+  # Work on copies to avoid mutating the original objects
+  expected_copy = expected.__class__()
+  expected_copy.CopyFrom(expected)
+  actual_copy = actual.__class__()
+  actual_copy.CopyFrom(actual)
+
+  try:
+    normalize_tensor_protos(expected_copy)
+    normalize_tensor_protos(actual_copy)
+  except ValueError as e:
+    test_case.fail(f"Proto normalization failed: {e}. {msg}")
+
+  compare.assertProtoEqual(
+      test_case, expected_copy, actual_copy, msg=msg, **kwargs
+  )
+
+
 class MeridianTestCase(parameterized.TestCase):
   """Base test class for Meridian providing backend-aware utilities.
 

meridian/constants.py

@@ -662,6 +662,7 @@ CURRENT_SPEND = 'current_spend'
 
 # Media summary metrics.
 SPEND = 'spend'
+COST = 'cost'
 IMPRESSIONS = 'impressions'
 ROI = 'roi'
 OPTIMIZED_ROI = 'optimized_roi'

meridian/data/load.py

@@ -35,6 +35,8 @@ __all__ = [
     'InputDataLoader',
     'XrDatasetDataLoader',
     'DataFrameDataLoader',
+    'CoordToColumns',
+    'CsvDataLoader',
 ]
 
 

meridian/model/eda/__init__.py

@@ -17,4 +17,3 @@
 from meridian.model.eda import eda_engine
 from meridian.model.eda import eda_outcome
 from meridian.model.eda import eda_spec
-from meridian.model.eda import meridian_eda

meridian/model/eda/constants.py

@@ -14,8 +14,18 @@
 
 """Constants specific to MeridianEDA."""
 
-# EDA Plotting properties
+# EDA Engine constants
+COST_PER_MEDIA_UNIT = 'cost_per_media_unit'
+
 VARIABLE_1 = 'var1'
 VARIABLE_2 = 'var2'
-VARIABLE = 'var'
 CORRELATION = 'correlation'
+ABS_CORRELATION_COL_NAME = 'abs_correlation'
+
+# EDA Plotting properties
+VARIABLE = 'var'
+VALUE = 'value'
+NATIONALIZE = 'nationalize'
+MEDIA_IMPRESSIONS_SCALED = 'media_impressions_scaled'
+IMPRESSION_SHARE_SCALED = 'impression_share_scaled'
+SPEND_SHARE = 'spend_share'