transformer_engine/jax/flax/transformer.py · d725686612d633c87d8845fba08d0fe5b7c7862a · OpenDAS / TransformerEngine

"projects/xcl-lens/src/xcl_lens/runner.py" did not exist on "6e26d46d4ddb1c192bb9c423dc55a3f231a6843f"

[JAX] Expose context parallel params to jax DPA api (#1292) · d7256866

Md Fahim Faysal Khan authored Nov 04, 2024



Exposed context parallel params to DPA api
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>

---------
Signed-off-by: Md Fahim Faysal Khan <mdfahimfaysa@nvidia.com>
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com>
Co-authored-by: Michael Goldfarb <mgoldfarb@nvidia.com>

d7256866

transformer.py 87.3 KB

Replace transformer.py