From 0fb638c260540a18d7648a59e7a4a813e9a2ea4b Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5@gmail.com>
Date: Tue, 18 Nov 2025 10:58:10 +0000
Subject: [PATCH] Reenabled examples tests

---
 .github/workflows/flax_test.yml    | 3 ---
 pyproject.toml                     | 8 ++++----
 tests/download_dataset_metadata.sh | 2 +-
 tests/run_all_tests.sh             | 3 +--
 4 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/flax_test.yml b/.github/workflows/flax_test.yml
index d37525493..00c09c4c0 100644
--- a/.github/workflows/flax_test.yml
+++ b/.github/workflows/flax_test.yml
@@ -118,9 +118,6 @@ jobs:
         else
           uv pip install "jax==${{ matrix.jax-version }}" "jaxlib==${{ matrix.jax-version }}"
         fi
-        if [[ "${{ matrix.test-type }}" == "pytest" ]]; then
-          uv pip install -U tensorflow-datasets
-        fi
     - name: Test with ${{ matrix.test-type }}
       run: |
         if [[ "${{ matrix.test-type }}" == "doctest" ]]; then
diff --git a/pyproject.toml b/pyproject.toml
index 4dc0f4dde..aefae59d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,10 +46,10 @@ testing = [
     "ml-collections",
     "mypy",
     "opencv-python",
-    # Set protobuf version for python 3.13+ to prevent error in
+    # Set protobuf version to prevent error in
     # examples/mnist/train_test.py::TrainTest::test_train_and_evaluate
     # Failed to construct dataset "mnist", builder_kwargs "{}": Value out of range: 11594722
-    "protobuf<6; python_version>='3.13'",
+    "protobuf<6",
     "pytest",
     "pytest-cov",
     "pytest-custom_exit_code",
@@ -59,8 +59,8 @@ testing = [
     "sentencepiece==0.2.0",  # Segfault bug in 0.2.1
     "tensorflow_text>=2.11.0; platform_system!='Darwin' and python_version < '3.13'",
     "tensorflow_datasets",
-    "tensorflow>=2.12.0; python_version<'3.13'", # to fix Numpy np.bool8 deprecation error
-    "tensorflow>=2.20.0; python_version>='3.13' and python_version<'3.14'",
+    "tensorflow_cpu>=2.12.0; python_version<'3.13'", # to fix Numpy np.bool8 deprecation error
+    "tensorflow_cpu>=2.20.0; python_version>='3.13' and python_version<'3.14'",
     # Temporary fix for https://github.com/google/flax/issues/5143
     "keras<3.13",
     "torch",
diff --git a/tests/download_dataset_metadata.sh b/tests/download_dataset_metadata.sh
index dbde4e248..50083ea46 100644
--- a/tests/download_dataset_metadata.sh
+++ b/tests/download_dataset_metadata.sh
@@ -16,7 +16,7 @@ if [ -d "../.tfds/metadata" ]; then
   echo 'TFDS metadata already exists.';
 else
   echo 'TFDS metadata does not exist. Downloading...';
-  git clone --branch v4.8.2 --depth 3 --filter=blob:none --sparse https://github.com/tensorflow/datasets/
+  git clone --branch v4.9.9 --depth 3 --filter=blob:none --sparse https://github.com/tensorflow/datasets/
   cd datasets
   git sparse-checkout set tensorflow_datasets/testing/metadata
   mkdir ../../.tfds
diff --git a/tests/run_all_tests.sh b/tests/run_all_tests.sh
index e74cd7e00..9a5e45496 100755
--- a/tests/run_all_tests.sh
+++ b/tests/run_all_tests.sh
@@ -128,8 +128,7 @@ if $RUN_PYTEST; then
     if [[ $egd == *"_"* ]]; then
       continue
     fi
-    # skiping examples until tfds issue is resolved
-    # pytest $egd
+    pytest -vvv $egd
   done
 fi