transformer_engine/jax/sharding.py · 7e45be73bb8d513abe8785ee078ac88719bcd9f1 · OpenDAS / TransformerEngine

jberchtold-nvidia authored Oct 02, 2025



Fix shard map issue
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

be7f43f1

sharding.py 13.5 KB