tests/jax/test_custom_call_compute.py · cae1c436027cc028ee49c83463141ba67f0adca0 · OpenDAS / TransformerEngine

[JAX] Helper to disable TE custom calls + disable GemmPrimitive for non-MXFP8 recipes. (#1962) · 2a293456

Phuong Nguyen authored Jul 23, 2025



* add manage_primitives() helper

* disable GEMM primitives for non-MXFP8 recipes

* implement the NVTE_JAX_CUSTOM_CALLS + deprecate NVTE_JAX_CUSTOM_CALLS_RE

* replace NVTE_JAX_CUSTOM_CALLS_RE with NVTE_JAX_CUSTOM_CALLS in TE tests and examples

* fix use_jax_gemm contextmanager
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

2a293456

test_custom_call_compute.py 57.4 KB

Replace test_custom_call_compute.py