Reenabled examples tests

.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

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",

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

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