[JAX] Generalizing Activation Primitives - #810

Merged

phu0ngng merged 14 commits into

NVIDIA:mainfrom

phu0ngng:act_prim

May 3, 2024

Commits on May 1, 2024

templated primitives and respective C++ functions
phu0ngng
committed
e7f211e
fixes for LayerNormMLP, tests in test_custom_compute all passed
phu0ngng
committed
4e0227d
added default arg for pybind get_workspace_size funcs
phu0ngng
committed
2bd3296
fixes for TestTransFormer with non-gated act tests
phu0ngng
committed
7a74dec
renamed gelu to act
phu0ngng
committed
896689a
improved enum implementation, avoid using magic numbers
phu0ngng
committed
ae06729
Exposed C++ ActivationEnum to python side
phu0ngng
committed
2fa8d93
Changed error messages
phu0ngng
committed
13ae4e0
changed conditional check on input shape for dbias_cast_transpose
phu0ngng
committed
bc1ab2f
changed dtype (tol) for bias grad tests
phu0ngng
committed
3b84cc4
fixes so that layer_norm_fp8_mlp can take bias = None
phu0ngng
committed
36876b4

Commits on May 2, 2024

rs conflicts with main
phu0ngng
committed
b47b3ad
Merge branch 'main' of github.com:NVIDIA/TransformerEngine into act_prim
phu0ngng
committed
62fc389
Set bias = None in flax modules
phu0ngng
committed
569baf1