Reenabled examples tests

Author: vfdev-5

.github/workflows/flax_test.yml

@@ -129,7 +129,6 @@ jobs:
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run --no-sync tests/run_all_tests.sh --only-doctest
         elif [[ "${{ matrix.test-type }}" == "pytest" ]]; then
-          uv pip install -U tensorflow-datasets
           # temporary: install jax nightly
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run --no-sync tests/run_all_tests.sh --only-pytest

pyproject.toml

@@ -47,10 +47,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",

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