tests/jax/test_distributed_layernorm_mlp.py · 643fb0a0e0b07d7cbfc103522e8f2240f101eecb · OpenDAS / TransformerEngine

Johannes Reifferscheid authored Apr 14, 2025



* Add experimental Shardy support.

Production use is not yet recommended.

---------
Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

6117b20c

test_distributed_layernorm_mlp.py 15.3 KB