transformer_engine/jax/fused_attn.py · 387397a2afeabfe52de019ea9f03f46be3f47f2d · OpenDAS / TransformerEngine

[JAX] Fix JAX distributed unit tests (#521) · ea43b18e

zlsh80826 authored Nov 21, 2023



* Remove assertion for NO_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix JAX distributed unit tests name
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>

ea43b18e

fused_attn.py 8.54 KB

Replace fused_attn.py