From 863fa165f265688390d487fbe77c609ccc85b2e2 Mon Sep 17 00:00:00 2001
From: tom denton <sdenton4@gmail.com>
Date: Mon, 11 Mar 2024 09:59:20 -0700
Subject: [PATCH] Handle tensorflow version pinning in setup.py

---
 setup.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/setup.py b/setup.py
index 2f8c2ad3d..f8dca4b61 100644
--- a/setup.py
+++ b/setup.py
@@ -170,11 +170,11 @@ def has_ext_modules(self):
     python_requires=">=3.7, <3.12",
     install_requires=install_requires,
     extras_require={
-        "tensorflow": [require],
-        "tensorflow-gpu": [require.replace("tensorflow", "tensorflow-gpu")],
-        "tensorflow-cpu": [require.replace("tensorflow", "tensorflow-cpu")],
-        "tensorflow-rocm": [require.replace("tensorflow", "tensorflow-rocm")],
-        "tensorflow-aarch64": [require.replace("tensorflow", "tensorflow-aarch64")],
+        "tensorflow >= 2.15, <2.16": [require],
+        "tensorflow-gpu == 2.15, <2.16": [require.replace("tensorflow", "tensorflow-gpu")],
+        "tensorflow-cpu == 2.15, <2.16": [require.replace("tensorflow", "tensorflow-cpu")],
+        "tensorflow-rocm == 2.15, <2.16": [require.replace("tensorflow", "tensorflow-rocm")],
+        "tensorflow-aarch64 == 2.15, <2.16": [require.replace("tensorflow", "tensorflow-aarch64")],
     },
     package_data={
         ".": ["*.so"],