examples/jax/mnist/test_single_gpu_mnist.py · 8ec01e5e77908cab9b8278d87776b2e5253a6a63 · OpenDAS / TransformerEngine

[JAX] Use FP8 tolerances in FP8 tests (#501) · 8ec01e5e

Tim Moon authored Nov 07, 2023



* Use FP8 tolerances in JAX FP8 tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Programmatically compute expected floating point error
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Loosen tolerance for MNIST test
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

8ec01e5e

test_single_gpu_mnist.py 10.1 KB

Replace test_single_gpu_mnist.py