transformer_engine/jax/quantize/scaling_modes.py · 3a298e6b194c6cc659f79533a49a13b26cd9f560 · OpenDAS / TransformerEngine

[JAX] TensorUsage + FP8 GEMM with all layouts handling on BW (#1844) · 3a298e6b

Phuong Nguyen authored Jun 18, 2025



* TensorUsage + FP8 GEMM with all layouts handling on BW
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

3a298e6b

scaling_modes.py 25.7 KB

Replace scaling_modes.py