tests/jax/test_custom_call_compute.py · 637faccb0a3f946a9bfb70e9d7e64c5fb251c856 · 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

test_custom_call_compute.py 55.9 KB

Replace test_custom_call_compute.py