[CI] Reenable example tests

[vfdev-5]

Example tests are skipped starting from this PR: #5090 due to a failure with TFDS:

• https://github.com/google/flax/actions/runs/19378708804/job/55452481246

We should reenable them and fix the issues.

1. ValueError: Failed to construct dataset "mnist", builder_kwargs "{}": Value out of range: 11594722 is due to protobuf version. protobuf with versions: 5.29.5, 6.30.0, 6.31.0 should work. CI has 6.33.1 and fails.

2. jax.device_put_sharded is deprecated and used in flax.jax_utils.prefetch_to_device in imagenet example

• https://github.com/google/flax/actions/runs/19470992916/job/55718044245?pr=5098