transformer_engine/jax/cpp_extensions/softmax.py · 063ef88d1af050b9076daaa5b92ed2d0018796fa · 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

softmax.py 29.2 KB