fixed static args

Signed-off-by: Alp Dener <[email protected]>
NVIDIA · Nov 21, 2024 · e44d5cf · e44d5cf
1 parent 8a610d0
commit e44d5cf
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py
@@ -3,9 +3,8 @@
 # See LICENSE for license information.
 import warnings
 import operator
-from functools import reduce
+from functools import reduce, partial
 from typing import Optional, Tuple
-from collections.abc import Iterable
 
 import jax
 import jax.numpy as jnp
@@ -69,7 +68,7 @@ class CollectiveGemmPrimitive(BasePrimitive):
     """
 
     name = "te_gemm"
-    impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15)
+    impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15, 16)
     multiple_results = True
     inner_primitive = None
     outer_primitive = None