jax
bf6c15b5 - update pmap to flatten correctly (was a perf bug)

Commit
6 years ago
update pmap to flatten correctly (was a perf bug) also temporarily avoid DeviceTuples in optimizer states
Author
Committer
Parents
Loading