transformer_engine/jax/flax/transformer.py · 24e4f955d9de38402585bd0322c6f1e5dabd96cf · OpenDAS / TransformerEngine

[JAX] Flax params initialization with weight_dtype (#1481) · 24e4f955

Phuong Nguyen authored Feb 13, 2025



* initialization with weight_dtype
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

24e4f955

transformer.py 88.9 KB

Replace transformer.py