tests/jax/test_custom_call_compute.py · 44740c6cc8716ac6fb141775a4ab478e9a4143ae · OpenDAS / TransformerEngine

[JAX] Set `precision=HIGHEST` for the ref_grouped_gemm impl in the unit test (#1967) · 2d4644b7

Phuong Nguyen authored Jul 18, 2025



* set precision=HIGHEST for the ref_grouped_gemm impl in the unit test
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


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

2d4644b7

test_custom_call_compute.py 57.3 KB

Replace test_custom_call_compute.py