Two implementations of ZeRO-1 optimizer sharding in JAX
-
Updated
Jun 11, 2023 - Python
#
pjit
Here are 3 public repositories matching this topic...
Tensor Parallelism with JAX + Shard Map
-
Updated
Sep 29, 2023 - Python
Comprehensive JAX implementation of neural networks and scientific computing. Features distributed training, physics-informed networks, custom autodiff, and advanced optimization. Production-ready code with numerical stability, multi-device parallelism, and research-grade implementations.
machine-learning research deep-learning optimization automatic-differentiation parallel-computing transformers bayesian-methods scientific-computing neural-networks high-performance-computing differential-equations gpu-computing pmap distributed-training jax mixed-precision physics-informed pjit tpu-computing
-
Updated
Sep 1, 2025 - Jupyter Notebook
Improve this page
Add a description, image, and links to the pjit topic page so that developers can more easily learn about it.
Add this topic to your repo
To associate your repository with the pjit topic, visit your repo's landing page and select "manage topics."