transformer_engine/jax/cpp_extensions.py · 6459fd853c02a3d1fcc45c1ade2d3e83cffbe66b · OpenDAS / TransformerEngine

[JAX] SwiGLU Implementation (#773) · f85553ea

Phuong Nguyen authored Apr 24, 2024



* Implemented swiglu and silu
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Renamed nvte-*silu to nvte-*swish + generalized GetDBiasDact functions
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

f85553ea

cpp_extensions.py 236 KB

Replace cpp_extensions.py