examples/jax/encoder/test_model_parallel_encoder.py · 07afda98585511ede420971ca8223a8ac2c07d1a · OpenDAS / TransformerEngine

[JAX] Support Flax sharding constraints (#1933) · c0c12e20

jberchtold-nvidia authored Jul 16, 2025



* Support flax sharding constraints
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Add warning for deprecated TE logical axes
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Update examples
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

c0c12e20

test_model_parallel_encoder.py 20.8 KB

Replace test_model_parallel_encoder.py