tests/jax/test_custom_call_compute.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

test_custom_call_compute.py 22.5 KB

Replace test_custom_call_compute.py