Update dependency from tensorflow-cpu to tensorflow.

google · Nov 4, 2024 · e6fb636 · e6fb636
1 parent c99465b
commit e6fb636
Show file tree

Hide file tree

Showing 9 changed files with 715 additions and 415 deletions.
diff --git a/BUILD b/BUILD
@@ -127,7 +127,7 @@ py_wheel(
     }),
     python_requires = ">=3.9",
     python_tag = "INTERPRETER",
-    requires = ["tensorflow-cpu==2.18.0"],  # See also: requirements.in.
+    requires = ["tensorflow[and-cuda]==2.18.0"],  # See also: requirements.in.
     # The summary is tailored for each python version because PyPI prevents
     # wheel uploads for different versions which have the same contents.
     # Changing the summary is sufficient to allow re-uploads.

diff --git a/requirements.in b/requirements.in
@@ -6,5 +6,5 @@
 #   https://github.com/tensorflow/tensorflow/pull/74090
 # When the TensorFlow PR is merged, builds for python 3.12 can be re-enabled.
 # Note this is not the case as of TensorFlow 2.18.0.
-tensorflow-cpu==2.18.0  # See also: top level BUILD, wheel target.
+tensorflow[and-cuda]==2.18.0  # See also: top level BUILD, wheel target.
 black==24.10.0
diff --git a/requirements_3_10.txt b/requirements_3_10.txt
diff --git a/requirements_3_11.txt b/requirements_3_11.txt
diff --git a/requirements_3_12.txt b/requirements_3_12.txt
diff --git a/requirements_3_9.txt b/requirements_3_9.txt
diff --git a/tf_shell/test/BUILD b/tf_shell/test/BUILD
@@ -8,7 +8,7 @@ py_test(
     ],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -21,7 +21,7 @@ py_test(
     ],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -35,7 +35,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -49,7 +49,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -63,7 +63,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -77,7 +77,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -91,7 +91,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -105,7 +105,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -119,7 +119,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -133,7 +133,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -147,7 +147,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -160,7 +160,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -173,7 +173,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -186,7 +186,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -199,7 +199,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -212,7 +212,7 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -225,6 +225,6 @@ py_test(
     imports = ["./"],
     deps = [
         "//tf_shell:tf_shell_lib",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
diff --git a/tf_shell_ml/test/BUILD b/tf_shell_ml/test/BUILD
@@ -7,7 +7,7 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -17,7 +17,7 @@ py_test(
     srcs = ["mnist_noenc_backprop_test.py"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -28,7 +28,7 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -38,7 +38,7 @@ py_test(
     srcs = ["dropout_test.py"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -48,7 +48,7 @@ py_test(
     srcs = ["embedding_test.py"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -58,7 +58,7 @@ py_test(
     srcs = ["conv2d_test.py"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -68,7 +68,7 @@ py_test(
     srcs = ["max_pool2d_test.py"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -79,7 +79,7 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -90,7 +90,7 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -101,7 +101,7 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -112,6 +112,6 @@ py_test(
     tags = ["exclusive"],
     deps = [
         "//tf_shell_ml",
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
diff --git a/third_party/tensorflow/BUILD b/third_party/tensorflow/BUILD
@@ -6,7 +6,7 @@ py_binary(
     srcs = [":extract_tf_includes.py"],
     srcs_version = "PY3",
     deps = [
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )
 
@@ -15,7 +15,7 @@ py_binary(
     srcs = [":extract_tf_lib.py"],
     srcs_version = "PY3",
     deps = [
-        requirement("tensorflow-cpu"),
+        requirement("tensorflow"),
     ],
 )