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

References

Author

mattjj

Committer

mattjj

Parents