transformer_engine/jax/cpp_extensions.py · aaf9354861c447eb5d53cb2907572490c0673861 · OpenDAS / TransformerEngine

[JAX] Allow multi-dims for dgamma and dbeta in LN descriptor. (#780) · aaf93548

Ming-Xu Huang authored Apr 19, 2024



* Allow multi-dims for dgamma and dbeta in LN descriptor.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the jit error in examples/jax
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>

aaf93548

cpp_extensions.py 184 KB

Replace cpp_extensions.py