transformer_engine/jax/cpp_extensions.py · 87e4d6c33ccd4242a27c3fadad53d604a3710396 · OpenDAS / TransformerEngine

[JAX] Fixes for the issue with ActLuPrimitive in PAXML (#837) · 87e4d6c3

Phuong Nguyen authored May 09, 2024



* fixes for ActLuPrimitive in PAXML

* changed indices for arg_infos in sharding func in dbias_cast_transpose primitive

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

87e4d6c3

cpp_extensions.py 183 KB

Replace cpp_extensions.py