diff --git a/MODULE.bazel b/MODULE.bazel
index a829273..a763404 100644
--- a/MODULE.bazel
+++ b/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")
diff --git a/beam/examples/requirements.txt b/beam/examples/requirements.txt
index d0cb7fa..f924f2d 100644
--- a/beam/examples/requirements.txt
+++ b/beam/examples/requirements.txt
@@ -1,3 +1 @@
 array-record[beam]
-google-cloud-storage==2.11.0
-tensorflow==2.14.0
\ No newline at end of file
diff --git a/cpp/array_record_writer.h b/cpp/array_record_writer.h
index 385313d..2e1d5c4 100644
--- a/cpp/array_record_writer.h
+++ b/cpp/array_record_writer.h
@@ -217,11 +217,11 @@ class ArrayRecordWriterBase : public riegeli::Object {
     // 31 in 64-bit build).
     //
     // Default: `absl::nullopt`.
-    Options& set_window_log(std::optional<int> window_log) {
+    Options& set_window_log(absl::optional<int> window_log) {
       compressor_options_.set_window_log(window_log);
       return *this;
     }
-    std::optional<int> window_log() const {
+    absl::optional<int> window_log() const {
       return compressor_options_.window_log();
     }
 
diff --git a/oss/build_whl.sh b/oss/build_whl.sh
index 6084176..1dace4d 100755
--- a/oss/build_whl.sh
+++ b/oss/build_whl.sh
@@ -87,12 +87,12 @@ 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
diff --git a/setup.py b/setup.py
index 79990c7..90df110 100644
--- a/setup.py
+++ b/setup.py
@@ -9,11 +9,19 @@
     '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 +42,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=[