Two implementations of ZeRO-1 optimizer sharding in JAX
-
Updated
Jun 11, 2023 - Python
#
pjit
Here are 4 public repositories matching this topic...
Tensor Parallelism with JAX + Shard Map
-
Updated
Sep 29, 2023 - Python
Numerics and Systems Lab in JAX: 21 notebooks executed end to end in CI, plus jax_nsl, a 275-test library covering numerically stable primitives, autodiff tools (Hessian-vector products, implicit differentiation), transforms, preconditioned linear algebra, training code and multi-device parallelism.
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
Nov 1, 2026 - Jupyter Notebook
P-JiT 21K for Stable Diffusion WebUI Forge Neo: Linum's small pixel-space text-to-image model, with no VAE, on 8 GB GPUs
-
Updated
Nov 7, 2026 - Python
Add this topic to your repo
To associate your repository with the pjit topic, visit your repo's landing page and select "manage topics."