Donate buffers for "optimizer", "batch" and "batch_stats" with donate_argnums to the pmapped computation to reduce memory requirements.
Donate buffers for "optimizer", "batch" and "batch_stats" with donate_argnums to the pmapped computation to reduce memory requirements.

PiperOrigin-RevId: 379415095

google/init2winit