transformer_engine/jax/cpp_extensions/activation.py · 57b4d7bc0350917cd2122b07f144bfeb4c04eb0b · OpenDAS / TransformerEngine

Phuong Nguyen authored Sep 22, 2025



* remove import jax.extend.ffi
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

57b4d7bc

activation.py 41.1 KB