flax

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

#388

Merged