Reenabled examples tests

.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

flax/nnx/variablelib.py

@@ -1712,25 +1712,25 @@ def __contains__(self, item) -> bool:
 
   def __eq__(self, other) -> bool:
     if isinstance(other, Variable):
-      other = other.value
-    return self.value.__eq__(other)  # type: ignore
+      other = other[...]
+    return self[...].__eq__(other)  # type: ignore
 
   def __iadd__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value += x` instead.'
+      'Use `variable[...] += x` instead.'
     )
 
   def __isub__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value -= x` instead.'
+      'Use `variable[...] -= x` instead.'
     )
 
   def __imul__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value *= x` instead.'
+      'Use `variable[...] *= x` instead.'
     )
 
   def __imatmul__(self: V, other) -> V:

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