From ae05b63ac2da906c657f688f81f4a49113dbba78 Mon Sep 17 00:00:00 2001
From: thomaswc <thomaswc@google.com>
Date: Thu, 28 Mar 2024 13:05:12 -0700
Subject: [PATCH] Fix AutoBNN's pip package by including all of its
 dependencies (and not just the ones it has in addition to tfp's).

PiperOrigin-RevId: 620032524
---
 spinoffs/autobnn/setup.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/spinoffs/autobnn/setup.py b/spinoffs/autobnn/setup.py
index 184fb98611..9e30e5126d 100644
--- a/spinoffs/autobnn/setup.py
+++ b/spinoffs/autobnn/setup.py
@@ -46,10 +46,13 @@
         'bayeux-ml',
         'chex',
         'flax',
+        'jax',
         'jaxtyping',
         'matplotlib',
+        'numpy',
         'pandas',
         'scipy',
+        'tensorflow-probability'
         ],
     # Add in any packaged data.
     exclude_package_data={'': ['BUILD']},