flax

e25e20c8 - [tfds] Fix tfds.as_numpy which now returns a reusable iterable.

[tfds] Fix tfds.as_numpy which now returns a reusable iterable. `next(ds)` -> `next(iter(ds))` Context: https://github.com/tensorflow/datasets/issues/2270 PiperOrigin-RevId: 326199391