google / flax

Flax is a neural network library for JAX that is designed for flexibility.

Home Page:https://flax.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Add HOWTO: Data-parallel training

marcvanzee opened this issue · comments

See #1982, this HOWTO will implement data-parallel training using pmap in the seq2seq example.

@gmittal mentioned they would like to implement this based on #1982

@gmittal any updates on this?

FYI I am closing this issue in favor of #2116

@marcvanzee sorry for delay -- this is still in-progress and will have a PR for this in the next week.