transformer_engine/jax/cpp_extensions.py · 07291027ed353287149e9df6030862e1e815f32f · OpenDAS / TransformerEngine

"git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "665c48963be11b2e5cb7209cd25f884129e5c284"

[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