tests/jax/test_praxis_layers.py · f85553ea369da15fd726ab279818e415be48a228 · 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_praxis_layers.py 51.9 KB

Replace test_praxis_layers.py