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

@@ -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

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

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(

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 = [],
 )