From b15d2b87c46f79eea314285cf02bf3095a725b2b Mon Sep 17 00:00:00 2001
From: Mihir Paradkar <mparadkar@google.com>
Date: Fri, 26 May 2023 10:05:38 -0700
Subject: [PATCH] Unpins packages that don't need pinning (protobuf needs to be
 pinned to the intersection of beam and TF requirements)

PiperOrigin-RevId: 535647770
---
 kokoro/github/ubuntu/cpu/build.sh           | 1 +
 kokoro/github/ubuntu/cpu/build_versioned.sh | 1 +
 setup.py                                    | 8 ++------
 tensorflow_gnn/experimental/sampler/BUILD   | 1 -
 4 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/kokoro/github/ubuntu/cpu/build.sh b/kokoro/github/ubuntu/cpu/build.sh
index da127534..30d7f962 100644
--- a/kokoro/github/ubuntu/cpu/build.sh
+++ b/kokoro/github/ubuntu/cpu/build.sh
@@ -37,6 +37,7 @@ ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn
 tag_filters="-no_oss,-oss_excluded"
 
 bazel clean
+pip install --upgrade pip
 pip install -r requirements-dev.txt --progress-bar off
 pip install tf-nightly --progress-bar off --upgrade
 # We need to remove the dependency on tensorflow to test nightly
diff --git a/kokoro/github/ubuntu/cpu/build_versioned.sh b/kokoro/github/ubuntu/cpu/build_versioned.sh
index 52389159..4f13bb90 100644
--- a/kokoro/github/ubuntu/cpu/build_versioned.sh
+++ b/kokoro/github/ubuntu/cpu/build_versioned.sh
@@ -44,6 +44,7 @@ ln -s "$(pwd)"/tensorflow_gnn "$TEST_ROOT"/tensorflow_gnn
 tag_filters="-no_oss,-oss_excluded${TAG_FILTERS}"
 
 bazel clean
+pip install --upgrade pip
 pip install -r requirements-dev.txt --progress-bar off
 pip install tensorflow=="${TF_VERSION}" --progress-bar off --upgrade
 python3 setup.py bdist_wheel
diff --git a/setup.py b/setup.py
index ce1c0d5a..417d769f 100644
--- a/setup.py
+++ b/setup.py
@@ -168,8 +168,6 @@ def get_version():
     ],
     namespace_packages=[],
     install_requires=[
-        # TODO(b/283835852): Unpin version when Vizier takes protobuf >= 4.0.
-        'apache-beam~=2.46.0',
         'google-vizier>=0.0.13',
         'ml-collections',
         'networkx',
@@ -178,10 +176,8 @@ def get_version():
         'tensorflow>=2.9.0; platform_machine != "arm64" or platform_system != "Darwin"',
         'tensorflow-macos>=2.9.0; platform_machine == "arm64" and platform_system == "Darwin"',
         # pylint:enable=g-line-too-long
-        # TODO(b/283835852): Remove when Vizier takes protobuf >= 4.0.
-        'protobuf~=3.20.3',
-        # TODO(b/283835852): Remove when Vizier takes protobuf >= 4.0.
-        'wrapt~=1.14.1',
+        'apache-beam',
+        'wrapt~=1.11.2',
     ],
     python_requires='>=3.8,<4',
     packages=find_namespace_packages(
diff --git a/tensorflow_gnn/experimental/sampler/BUILD b/tensorflow_gnn/experimental/sampler/BUILD
index 22040813..99eed118 100644
--- a/tensorflow_gnn/experimental/sampler/BUILD
+++ b/tensorflow_gnn/experimental/sampler/BUILD
@@ -19,7 +19,6 @@ py_proto_library(
 py_test(
     name = "core_test",
     srcs = ["core_test.py"],
-    tags = ["no_oss"],  # TODO(b/284063221)
     python_version = "PY3",
     deps = [],
 )