tests/jax/test_distributed_softmax.py · 6d178b4e0465cb6739d7213a6814325aa92d2e38 · OpenDAS / TransformerEngine

[JAX] Reduce L1 tests/jax/test_distributed_softmax.py test runtime (#2031) · 6d178b4e

jberchtold-nvidia authored Aug 06, 2025



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

* Reduce softmax test shape sizes
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Switch softmax tests to use shardy by default
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

6d178b4e

test_distributed_softmax.py 7.14 KB

Replace test_distributed_softmax.py