NVIDIA · tdophung · Feb 25, 2026 · Feb 26, 2026 · Feb 26, 2026 · Feb 28, 2026
diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py
diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py
diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py
@@ -9,3 +9,4 @@
 from .quantization import *
 from .softmax import *
 from .gemm import *
+from .router import *