tests/jax/test_layer.py · 5d937c57c6b8f964fe31e48ec48842acbe161743 · OpenDAS / TransformerEngine

Ming-Xu Huang authored Apr 08, 2023



* Rename enable_fp8 to is_fp8_enabled.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding an API to get an instance of  DelayedScaling which is set via fp8_autocast.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>

5d937c57

test_layer.py 25 KB

Replace test_layer.py