tests/jax/test_custom_call_compute.py · f8c2af4cbeea2b505b1010623080a8cf8a179229 · OpenDAS / TransformerEngine

[JAX] Fix failing L2 JAX unit tests (#1735) · fe31af80

jberchtold-nvidia authored May 06, 2025



* Fix L2 test_custom_call_compute.py L2 tests
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Fix test_helper.py
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

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

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

fe31af80

test_custom_call_compute.py 53.8 KB

Replace test_custom_call_compute.py