
Train models with JAX and Flax.
Language models, diffusion and JEPA, trained by one trainer on one device or a mesh.
uv pip install "dew-ml @ git+https://github.com/AshishKumar4/dew"Image and video diffusion
DiT, U-Net and MMDiT backbones, 18 solvers, classifier-free guidance.

Language models
Pretraining, chat fine-tuning and cached generation on one trainer.

One mesh
MeshSpec(fsdp=4, tensor=2) runs the same code on eight devices.

Post-training
SFT, DPO, GRPO and PPO continue the weights you trained.
