paxml (pypi): Framework to configure and run machine learning experiments on top of Jax.
Keywords: c4|gpt|jax|large-language-models|llm|model-flops|parallelism