flax

71e4432d - Adds `config.shuffle_buffer_size` and `config.prefetch`.

Adds `config.shuffle_buffer_size` and `config.prefetch`. This allows for more fine-grained control, which is particularly useful in a memory-constrained setup (e.g. public Colab runtimes). PiperOrigin-RevId: 520615589