| Acme: A research framework for distributed reinforcement learning MW Hoffman, B Shahriari, J Aslanides, G Barth-Maron, N Momchev, ... arXiv preprint arXiv:2006.00979, 2020 | 232 | 2020 |
| Haiku: Sonnet for JAX T Hennigan, T Cai, T Norman, I Babuschkin Version 0.0 3, 147, 2020 | 163* | 2020 |
| The DeepMind JAX Ecosystem I Babuschkin, K Baumli, A Bell, S Bhupatiraju, J Bruce, P Buchlovsky, ... URL http://github. com/deepmind 24, 25, 2020 | 140* | 2020 |
| Synthesizing optimal parallelism placement and reduction strategies on hierarchical systems for deep learning N Xie, T Norman, D Grewe, D Vytiniotis Proceedings of Machine Learning and Systems 4, 548-566, 2022 | 12 | 2022 |
| The DeepMind JAX Ecosystem, 2020 IB DeepMind, K Baumli, A Bell, S Bhupatiraju, J Bruce, P Buchlovsky, ... URL http://github. com/google-deepmind, 0 | 11 | |
| Automap: Towards ergonomic automated parallelism for ML models M Schaarschmidt, D Grewe, D Vytiniotis, A Paszke, GS Schmid, T Norman, ... arXiv preprint arXiv:2112.02958, 2021 | 9 | 2021 |
| PartIR: Composing SPMD Partitioning Strategies for Machine Learning S Alabed, B Chrzaszcz, J Franco, D Grewe, D Maclaurin, J Molloy, ... arXiv preprint arXiv:2401.11202, 2024 | | 2024 |
| Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR S Alabed, D Grewe, J Franco, B Chrzaszcz, T Natan, T Norman, NA Rink, ... arXiv preprint arXiv:2210.06352, 2022 | | 2022 |
