tests/jax/test_custom_call_compute.py · 40a30a5f8de9d97105a31bad7f23ed40abdf739b · OpenDAS / TransformerEngine

[JAX] GroupedDense v.2 without dynamic shape (#1875) · c9d7f3f2

Phuong Nguyen authored Jun 12, 2025



* Implemented GroupedDense and TestGroupedDense for BF16, FP16, and FP8 
* Fix GroupedGemmFFI cuBLAS workspace alignment bug
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>

c9d7f3f2

test_custom_call_compute.py 54.5 KB

Replace test_custom_call_compute.py