Add [and-cuda] optional dependencies to pyproject.toml.

PiperOrigin-RevId: 723585524

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+* Add `[and-cuda]` optional dependencies to install correct version of
+  `tensorflow`.
 * Add `non_media_baseline_values` argument to `summary_metrics`,
   `baseline_summary_metrics` and `expected_vs_actual` methods.
 

README.md

@@ -37,13 +37,29 @@ Note: This project has been tested on V100 and T4 GPU using 16 GB of RAM.
 To install Meridian, run the following command to automatically install the
 latest release from PyPI.
 
+For GPU users:
+
+```sh
+$ pip install --upgrade google-meridian[and-cuda]
+```
+
+For CPU users:
+
 ```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
 ```

demo/Meridian_Getting_Started.ipynb

@@ -87,13 +87,13 @@
       "outputs": [],
       "source": [
         "# Install meridian: from PyPI @ latest release\n",
-        "!pip install --upgrade google-meridian[colab]\n",
+        "!pip install --upgrade google-meridian[colab,and-cuda]\n",
         "\n",
         "# Install meridian: from PyPI @ specific version\n",
-        "#!pip install google-meridian==1.0.0\n",
+        "# !pip install google-meridian[colab,and-cuda]==1.0.1\n",
         "\n",
         "# Install meridian: from GitHub @HEAD\n",
-        "#!pip install --upgrade \"google-meridian[colab] @ git+https://github.com/google/meridian.git\""
+        "# !pip install --upgrade \"google-meridian[colab,and-cuda] @ git+https://github.com/google/meridian.git\""
       ]
     },
     {

pyproject.toml

@@ -62,6 +62,11 @@ dev = [
 colab = [
   "psutil",
 ]
+# GPU deps
+# Installed through `pip install -e .[and-cuda]`
+and-cuda = [
+  "tensorflow[and-cuda] >= 2.16, < 2.17",
+]
 
 [tool.pyink]
 # Formatting configuration to follow Google style-guide