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

google_meridian-1.3.1.dist-info/METADATA

@@ -0,0 +1,209 @@
+Metadata-Version: 2.4
+Name: google-meridian
+Version: 1.3.1
+Summary: Google's open source mixed marketing model library, helps you understand your return on investment and direct your ad spend with confidence. 
+Author-email: The Meridian Authors <no-reply@google.com>
+Project-URL: homepage, https://github.com/google/meridian
+Project-URL: repository, https://github.com/google/meridian
+Project-URL: changelog, https://github.com/google/meridian/blob/main/CHANGELOG.md
+Project-URL: documentation, https://developers.google.com/meridian
+Keywords: mmm
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Other/Nonlisted Topic
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: arviz
+Requires-Dist: altair>=5
+Requires-Dist: immutabledict
+Requires-Dist: joblib
+Requires-Dist: natsort<8,>=7.1.1
+Requires-Dist: numpy<3,>=2.0.2
+Requires-Dist: pandas<3,>=2.2.2
+Requires-Dist: patsy<1,>=0.5.3
+Requires-Dist: scipy<2,>=1.13.1
+Requires-Dist: statsmodels>=0.14.5
+Requires-Dist: tensorflow<2.19,>=2.18
+Requires-Dist: tensorflow-probability<0.26,>=0.25
+Requires-Dist: tf-keras<2.19,>=2.18
+Requires-Dist: xarray
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Requires-Dist: pylint>=2.6.0; extra == "dev"
+Requires-Dist: pyink; extra == "dev"
+Provides-Extra: colab
+Requires-Dist: psutil; extra == "colab"
+Requires-Dist: python-calamine; extra == "colab"
+Provides-Extra: and-cuda
+Requires-Dist: tensorflow[and-cuda]<2.19,>=2.18; extra == "and-cuda"
+Provides-Extra: mlflow
+Requires-Dist: mlflow; extra == "mlflow"
+Provides-Extra: jax
+Requires-Dist: jax==0.4.26; extra == "jax"
+Requires-Dist: jaxlib==0.4.26; extra == "jax"
+Requires-Dist: tensorflow-probability[substrates-jax]==0.25.0; extra == "jax"
+Provides-Extra: schema
+Requires-Dist: mmm-proto-schema; extra == "schema"
+Requires-Dist: semver; extra == "schema"
+Dynamic: license-file
+
+# About Meridian
+
+Marketing mix modeling (MMM) is a statistical analysis technique that measures
+the impact of marketing campaigns and activities to guide budget planning
+decisions and improve overall media effectiveness. MMM uses aggregated data to
+measure impact across marketing channels and account for non-marketing factors
+that impact sales and other key performance indicators (KPIs). MMM is
+privacy-safe and does not use any cookie or user-level information.
+
+Meridian is an MMM framework that enables advertisers to set up and run their
+own in-house models. Meridian helps you answer key questions such as:
+
+*   How did the marketing channels drive my revenue or other KPI?
+*   What was my marketing return on investment (ROI)?
+*   How do I optimize my marketing budget allocation for the future?
+
+Meridian is a highly customizable modeling framework that is based on
+[Bayesian causal inference](https://developers.google.com/meridian/docs/causal-inference/bayesian-inference).
+It is capable of handling large scale geo-level data, which is encouraged if
+available, but it can also be used for national-level modeling. Meridian
+provides clear insights and visualizations to inform business decisions around
+marketing budget and planning. Additionally, Meridian provides methodologies to
+support calibration of MMM with experiments and other prior information, and to
+optimize target ad frequency by utilizing reach and frequency data.
+
+If you are using LightweightMMM, see the
+[migration guide](https://developers.google.com/meridian/docs/migrate) to help
+you understand the differences between these MMM projects.
+
+## Install Meridian
+
+Python 3.11 or 3.12 is required to use Meridian. We also recommend using a
+minimum of 1 GPU.
+
+Note: This project has been tested on T4 GPU using 16 GB of RAM.
+
+To install Meridian, run the following command to automatically install the
+latest release from PyPI.
+
+*   For Linux-GPU users:
+
+    Note: CUDA toolchain and a compatible GPU device is necessary for
+    `[and-cuda]` extra to activate.
+
+    ```sh
+    $ pip install --upgrade google-meridian[and-cuda]
+    ```
+
+*   For macOS and general CPU users:
+
+    Note: There is no official GPU support for macOS.
+
+    ```sh
+    $ pip install --upgrade google-meridian
+    ```
+
+Alternatively, run the following command to install the most recent, unreleased
+version from GitHub.
+
+*   For GPU users:
+
+    ```sh
+    $ pip install --upgrade "google-meridian[and-cuda] @ git+https://github.com/google/meridian.git"
+    ```
+
+*   For CPU users:
+
+    ```sh
+    $ pip install --upgrade git+https://github.com/google/meridian.git
+    ```
+
+We recommend to install Meridian in a fresh
+[virtual environment](https://virtualenv.pypa.io/en/latest/user_guide.html#quick-start)
+to make sure that correct versions of all the dependencies are installed, as
+defined in [pyproject.toml](https://github.com/google/meridian/blob/main/pyproject.toml).
+
+## How to use the Meridian library
+
+To get started with Meridian, you can run the code programmatically using sample
+data with the [Getting Started Colab][3].
+
+The Meridian model uses a holistic MCMC sampling approach called
+[No U Turn Sampler (NUTS)](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/mcmc/NoUTurnSampler)
+which can be compute intensive. To help with this, GPU support has been
+developed across the library (out-of-the-box) using tensors. We recommend
+running your Meridian model on GPUs to get real time optimization results and
+significantly reduce training time.
+
+## Meridian Documentation & Tutorials
+
+The following documentation, colab, and video resources will help you get
+started quickly with using Meridian:
+
+| Resource                    | Description                                    |
+| --------------------------- | ---------------------------------------------- |
+| [Meridian documentation][1] | Main landing page for Meridian documentation.  |
+| [Meridian basics][2]        | Learn about Meridian features, methodologies, and the model math. |
+| [Getting started colab][3]  | Install and quickly learn how to use Meridian with this colab tutorial using sample data. |
+| [User guide][4]             | A detailed walk-through of how to use Meridian and generating visualizations using your own data. |
+| [Pre-modeling][5]           | Prepare and analyze your data before modeling. |
+| [Modeling][6]               | Modeling guidance for model refinement and edge cases. |
+| [Post-modeling][7]          | Post-modeling guidance for model fit, visualizations, optimizations, refreshing the model, and debugging. |
+| [Migrate from LMMM][8]      | Learn about the differences between Meridian and LightweightMMM as you consider migrating. |
+| [API Reference][9]          | API reference documentation for the Meridian package. |
+| [Reference list][10]        | White papers and other referenced material.    |
+
+[1]: https://developers.google.com/meridian
+[2]: https://developers.google.com/meridian/docs/basics/meridian-introduction
+[3]: https://developers.google.com/meridian/notebook/meridian-getting-started
+[4]: https://developers.google.com/meridian/docs/user-guide/installing
+[5]: https://developers.google.com/meridian/docs/pre-modeling/collect-data
+[6]: https://developers.google.com/meridian/docs/advanced-modeling/control-variables
+[7]: https://developers.google.com/meridian/docs/post-modeling/model-fit
+[8]: https://developers.google.com/meridian/docs/migrate
+[9]: https://developers.google.com/meridian/reference/api/meridian
+[10]: https://developers.google.com/meridian/docs/reference-list
+
+## Support
+
+**Questions about methodology**: Please see the [Modeling](https://developers.google.com/meridian/docs/basics/meridian-introduction) tab in the technical documentation.
+
+**Issues installing or using Meridian**: Feel free to post questions in the
+[Discussions](https://github.com/google/meridian/discussions) or [Issues](https://github.com/google/meridian/issues) tabs of the Meridian GitHub repository. The Meridian team responds to
+these questions weekly in batches, so please be patient and don't reach out
+directly to your Google Account teams.
+
+**Bug reports**: Please post bug reports to the [Issues](https://github.com/google/meridian/issues)
+tab of the Meridian GitHub repository. We also encourage the community to share
+tips and advice with each other on the [Issues](https://github.com/google/meridian/issues)
+tab. When our team addresses or resolves a new bug, we will notify you through
+the comments on the issue.
+
+**Feature requests**: Please post these to the [Discussions](https://github.com/google/meridian/discussions)
+tab of the Meridian GitHub repository. We have an internal roadmap for Meridian
+development, but would love your inputs for new feature requests so that we can
+prioritize them based on the roadmap.
+
+**Pull requests**: These are appreciated but are very difficult for us to merge
+because the code in this repository is linked to Google internal systems and has
+to pass internal review. If you submit a pull request and we believe that we can
+incorporate a change in the base code, we will reach out to you directly about
+this.
+
+## Citing Meridian
+
+To cite this repository:
+
+<!-- mdlint off(SNIPPET_INVALID_LANGUAGE) -->
+```BibTeX
+@software{meridian_github,
+  author = {Google Meridian Marketing Mix Modeling Team},
+  title = {Meridian: Marketing Mix Modeling},
+  url = {https://github.com/google/meridian},
+  version = {1.3.1},
+  year = {2025},
+}
+```

{google_meridian-1.3.0.dist-info → google_meridian-1.3.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
-google_meridian-1.3.0.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+google_meridian-1.3.1.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
 meridian/__init__.py,sha256=0fOT5oNZF7-pbiWWGUefV-ysafttieG079m1ijMFQO8,861
 meridian/constants.py,sha256=ZmMIoJDFQvKIOVG9oPOQ7Cj16wt4HDS5fCPBrz_KiLE,20308
-meridian/version.py,sha256=PmfWvor_VonhwnCn2anG7FaY2z7fcdLbE_3OeBHekCg,644
+meridian/version.py,sha256=mmCkrGRWB8mI33apVnverT2ysfJofSyhlXSCVNotj9U,644
 meridian/analysis/__init__.py,sha256=NWDhtRkKs-n66C6746rXT7Wk8tLdkrT9NDvuPG_B_5c,874
 meridian/analysis/analyzer.py,sha256=EUKoN69EEz4vt4BpyRvrwEoHt7eLZ8UM6emNe14Tbj0,219975
 meridian/analysis/formatter.py,sha256=AN2M4jdUV88XjNdi5-dK_mITES1J1Dk3Zs7DYo3OTKg,7290
@@ -25,9 +25,9 @@ meridian/analysis/templates/style.css,sha256=RODTWc2pXcG9zW3q9SEJpVXgeD-WwQgzLpm
 meridian/analysis/templates/style.scss,sha256=nSrZOpcIrVyiL4eC9jLUlxIZtAKZ0Rt8pwfk4H1nMrs,5076
 meridian/analysis/templates/summary.html.jinja,sha256=LuENVDHYIpNo4pzloYaCR2K9XN1Ow6_9oQOcOwD9nGg,1707
 meridian/analysis/templates/table.html.jinja,sha256=mvLMZx92RcD2JAS2w2eZtfYG-6WdfwYVo7pM8TbHp4g,1176
-meridian/backend/__init__.py,sha256=EzsRE_E3tmxNkRPuNannz6Ro0IddvEXaNZ365GG5X8o,34341
+meridian/backend/__init__.py,sha256=ftXcdb3tIky_m8exhD9RRhaEqTEvMcqZ8lGkkR1PjsE,39495
 meridian/backend/config.py,sha256=B9VQnhBfg9RW04GNbt7F5uCugByenoJzt-keFLLYEp8,3561
-meridian/backend/test_utils.py,sha256=XyYZ61o7fRCQq6rX75eON9AudvPHHdWzxHYK8_dJJvE,7753
+meridian/backend/test_utils.py,sha256=DYU5IpWiyM26aTY9Q84mOGRw0dQ9XmsZRsAJNFuZDp0,11667
 meridian/data/__init__.py,sha256=StIe-wfYnnbfUbKtZHwnAQcRQUS8XCZk_PCaEzw90Ww,929
 meridian/data/arg_builder.py,sha256=Kqlt88bOqFj6D3xNwvWo4MBwNwcDFHzd-wMfEOmLoPU,3741
 meridian/data/data_frame_input_data_builder.py,sha256=_hexZMFAuAowgo6FaOGElHSFHqhGnHQwEEBcwnT3zUE,27295
@@ -44,7 +44,7 @@ meridian/model/adstock_hill.py,sha256=HoRKjyL03pCTBz6Utof9wEvlQCFM43BvrEW_oupj7N
 meridian/model/knots.py,sha256=87kw5oa3T1k9GgT_aWXTqQx5XCxLsS2w1hnzc581XL0,26677
 meridian/model/media.py,sha256=skjy4Vd8LfDQWlqR_2lJ1qbG9UcS1dow5W45BAu4qk8,14599
 meridian/model/model.py,sha256=jMtfl7woWtJ8M8AX42QeZ5hUS8hlhPdZ-9OU8KahjKA,68984
-meridian/model/model_test_data.py,sha256=s8G1NjAYcF0-nCOPBZp6rxnKMLGdGo97Bg14mKumStE,23359
+meridian/model/model_test_data.py,sha256=XGBz8RGdCsjAUOmgxX3CfWSj-_hdq2Lc8saFCqmImwM,23901
 meridian/model/posterior_sampler.py,sha256=f3MayglIgBeBjWeXJU_RgT9cCugcjJ3aEjHqaWPsTbg,26806
 meridian/model/prior_distribution.py,sha256=ZArW4uXIPPQL6hRWiGZUzcHktbkjE_vOklvlbp9LR64,57662
 meridian/model/prior_sampler.py,sha256=iLvCefhA4WY0ENcnLK9471WUZPPyzQ1je58MRjxKv74,25460
@@ -52,11 +52,25 @@ meridian/model/spec.py,sha256=VlK6WJiPo2lzOF0O2judtJ6O3uEw7wYL5AT8bioq4gE,19188
 meridian/model/transformers.py,sha256=HxlVJitxP-wu-NOHU0tArFUZ4NAO3c7adAYj4Zvqnvo,8363
 meridian/model/eda/__init__.py,sha256=w3p7ZUZLq5TOEHm8n2P1CWjGrzuNrkqSSnVFdlw17Dk,812
 meridian/model/eda/constants.py,sha256=V9aOHQDvB3WAEyT0NE4gE8rqbStaGVh3XDlBPOKpuLc,739
-meridian/model/eda/eda_engine.py,sha256=5LCOtRSCLZf2ch70SN177UZl00VUWoswS_9Px3xUD9E,63672
+meridian/model/eda/eda_engine.py,sha256=5Ikgiz-6d3uTBan71WwsgPHOqk0fAy390ZLCq9L6HoY,64846
 meridian/model/eda/eda_outcome.py,sha256=P-0kNIbNXcyqMaNvFxiL3x6fhtYOL2trw-zPPZGXh5w,5670
 meridian/model/eda/eda_spec.py,sha256=diieYyZH0ee3ZLy0rGFMcWrrgiUrz2HctMwOrmtJR6w,2871
 meridian/model/eda/meridian_eda.py,sha256=GTdBaAtfsHS5s6P5ZESaeh1ElKV_o7dSqQosiMnFBKg,7537
-google_meridian-1.3.0.dist-info/METADATA,sha256=FHLqp_ZLBXWJHpZWW2gdWg8Fy37cUC7qAlsv_-cIrc8,22470
-google_meridian-1.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-google_meridian-1.3.0.dist-info/top_level.txt,sha256=nwaCebZvvU34EopTKZsjK0OMTFjVnkf4FfnBN_TAc0g,9
-google_meridian-1.3.0.dist-info/RECORD,,
+schema/__init__.py,sha256=Df2XKjMKa0ry7CTlDuJnyxuTdSgAd5za922wMlE98cg,681
+schema/serde/__init__.py,sha256=xyydIcWB5IUpcn3wu1m9HL1fK4gMWURbwTyRsQtolF0,975
+schema/serde/constants.py,sha256=aYtD_RuA0GCkpC4TIQq3VjMqEc837Wn-TlJNm-yn_4Y,1842
+schema/serde/distribution.py,sha256=jy3h6JD1TSs4gwociMis814sz_Fm2kFQ2UbkgjYJW9k,19347
+schema/serde/eda_spec.py,sha256=uOqBeZpUU3Dzzc19rU1LjHWmUhRmVcx8oIZvZfVJHT8,7180
+schema/serde/function_registry.py,sha256=GbgC5_9NDcA9Y7nqmdJ-4-LK5JPhhfI50Lmfy5ZBJOg,4858
+schema/serde/hyperparameters.py,sha256=Igm-PZmIozrsKZH6c-XkrU_Nlf8OAuxpnJJfv7W1SfQ,13524
+schema/serde/inference_data.py,sha256=DrwE9hU8LMrl0z8W_sUSIaPrRdym_lu0iOqpT4KZxsA,3623
+schema/serde/marketing_data.py,sha256=yb-fRTe84Sjg7-v3wsvYRRXvrxLSFWSenO0_ikMvUpk,44845
+schema/serde/meridian_serde.py,sha256=ZG05JaBG4LW8mhl-Cunje9Q6xyR4tyNTtLYedzMBYjA,15985
+schema/serde/serde.py,sha256=8vUqhJxvZgX9UY3rXTyWJznRgapwDzzaHXDHwV_kKTA,1612
+schema/serde/test_data.py,sha256=7hfEWyvZ9WcAkVAOXt6elX8stJlsfhfd-ASlHo9SRb8,107342
+schema/utils/__init__.py,sha256=AkC4NMbmXC3PFBY9dFYxlf3qFsxt5OOBVdc9zmFXsC8,675
+schema/utils/time_record.py,sha256=-KzHFjvSBUUXsfESPAfcJP_VFxaFLqj90Ac0kgKWfpI,4624
+google_meridian-1.3.1.dist-info/METADATA,sha256=wu5D6r6v46vd5g4uKLfXEVuhxkO3jwEgS68wn6m0jR4,9547
+google_meridian-1.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+google_meridian-1.3.1.dist-info/top_level.txt,sha256=yWkWDLV_UUanhKmk_xNPiKNdPDl1oyU1sBYwEnhaSf4,16
+google_meridian-1.3.1.dist-info/RECORD,,

{google_meridian-1.3.0.dist-info → google_meridian-1.3.1.dist-info}/top_level.txt

@@ -1 +1,2 @@
 meridian
+schema

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."""
@@ -497,12 +613,23 @@ def _jax_convert_to_tensor(data, dtype=None):
   # 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, (list, np.ndarray)) and np.array(
+      data
+  ).dtype.kind in ("S", "U")
+
+  if is_string_target or (dtype is None and is_string_data):
+    return np.array(data, dtype=dtype)
+
   return jax_ops.asarray(data, dtype=dtype)
 
 
@@ -535,18 +662,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
+
+  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(*args, **kwargs):  # pylint: disable=unused-argument
+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 +929,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 +1061,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

@@ -15,11 +15,21 @@
 """Common testing utilities for Meridian, designed to be backend-agnostic."""
 
 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.
@@ -131,6 +141,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.