[JAX] Expert Parallelism: JAX primitives + VJPs#3036

phu0ngng wants to merge 22 commits into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

phu0ngng:phuong/ep-3-jaxphu0ngng/TransformerEngine:phuong/ep-3-jaxCopy head branch name to clipboard