examples/jax/encoder/test_single_gpu_encoder.py · a3e4e611cfbd6906adbd6bcc75b6717f28682365 · OpenDAS / TransformerEngine

"vscode:/vscode.git/clone" did not exist on "88a94cc50d057dbf9b4ecb0bc39159a225f5f780"

[JAX] Fully remove attn_type and set self_attn_mask_type default to 'causal' (#324) · a3e4e611

zlsh80826 authored Jul 18, 2023



* Fully remove attn_type and set self_attn_mask_type default to 'causal'
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix tests with new arguments
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Explicit self_attn_mask_type for examples
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>

* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

a3e4e611

test_single_gpu_encoder.py 11.3 KB

Replace test_single_gpu_encoder.py