Upgrade protobuf to match with TF 2.20.0rc0

PiperOrigin-RevId: 791214571

Author: mtsokol

MODULE.bazel

@@ -22,11 +22,11 @@ module(
 )
 
 bazel_dep(name = "rules_proto", version = "7.0.2")
-bazel_dep(name = "rules_python", version = "0.37.0")
-bazel_dep(name = "platforms", version = "0.0.10")
-bazel_dep(name = "protobuf", version = "28.3")
+bazel_dep(name = "rules_python", version = "1.0.0")
+bazel_dep(name = "platforms", version = "0.0.11")
+bazel_dep(name = "protobuf", version = "31.1")
 bazel_dep(name = "googletest", version = "1.15.2")
-bazel_dep(name = "abseil-cpp", version = "20240722.0")
+bazel_dep(name = "abseil-cpp", version = "20250127.0")
 bazel_dep(name = "abseil-py", version = "2.1.0")
 bazel_dep(name = "eigen", version = "3.4.0.bcr.3")
 bazel_dep(name = "riegeli", version = "0.0.0-20241218-3385e3c")

beam/examples/requirements.txt

@@ -1,3 +1 @@
 array-record[beam]
-google-cloud-storage==2.11.0
-tensorflow==2.14.0

oss/build_whl.sh

@@ -87,15 +87,15 @@ function main() {
   echo $(date) : "=== Output wheel file is in: ${DEST}"
 
   # Install ArrayRecord from the wheel and run smoke tests.
-  $PYTHON_BIN -m pip install --find-links="${DEST}" --pre array-record
+  $PYTHON_BIN -m pip install ${OUTPUT_DIR}/all_dist/array_record*.whl
   $PYTHON_BIN -c 'import array_record'
   $PYTHON_BIN -c 'from array_record.python import array_record_data_source'
   # TF is not available on Python 3.13 and above.
   if [ "$(uname)" != "Darwin" ] && (( "${PYTHON_MINOR_VERSION}" < 13 )); then
-    $PYTHON_BIN -m pip install jax tensorflow grain
+    $PYTHON_BIN -m pip install jax tensorflow>=2.20.0rc0 grain
     $PYTHON_BIN oss/test_import_grain.py
     $PYTHON_BIN oss/test_import_tensorflow.py
   fi
 }
 
-main
+main

setup.py

@@ -9,11 +9,17 @@
     'etils[epath]',
 ]
 
+TF_PACKAGE = ['tensorflow>=2.20.0rc0']
+
 BEAM_EXTRAS = [
-    'apache-beam[gcp]==2.53.0',
+    'apache-beam[gcp]>=2.53.0',
     'google-cloud-storage>=2.11.0',
-    'tensorflow>=2.14.0'
-]
+] + TF_PACKAGE
+
+TEST_EXTRAS = [
+    'jax',
+    'grain',
+] + TF_PACKAGE
 
 
 class BinaryDistribution(Distribution):
@@ -34,7 +40,7 @@ def has_ext_modules(self):
     package_data={'': ['*.so']},
     python_requires='>=3.10',
     install_requires=REQUIRED_PACKAGES,
-    extras_require={'beam': BEAM_EXTRAS},
+    extras_require={'beam': BEAM_EXTRAS, 'test': TEST_EXTRAS},
     url='https://github.com/google/array_record',
     license='Apache-2.0',
     classifiers=[