transformer_engine/jax/cpp_extensions/attention.py · 0e13788366b4a9cd21eb42df6fafcd11f6c4bd0a · OpenDAS / TransformerEngine

Reese Wang authored Mar 13, 2025



Make ffi compatible with jax 0.4
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

0e137883

attention.py 101 KB