transformer_engine/jax/cpp_extensions/gemm.py · 3fc1e4bf8d46850f215698b39b5625310194bab9 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Fix for TE GEMM - Always AllGather RHS non-contracting dims with FSDP axis (#2075) · 3fc1e4bf
Phuong Nguyen authored Aug 18, 2025
```
* fix fsdp
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
```
3fc1e4bf

gemm.py 52.2 KB