updated FWD/BWD wrappers for non-FP8 and FP8 gemm

Signed-off-by: Alp Dener <[email protected]>
NVIDIA · Dec 5, 2024 · b1b51c3 · b1b51c3
1 parent 13a8cd4
commit b1b51c3
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 117 deletions.
diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py
@@ -1074,6 +1074,7 @@ def fp8_gemm_impl(
     bias: Optional[ArrayLike] = None,
     gelu_input: Optional[ArrayLike] = None,
     out: Optional[ArrayLike] = None,
+    extra_out: Optional[ArrayLike] = None,
     out_amax: Optional[ArrayLike] = None,
     out_scale: Optional[ArrayLike] = None,
     out_dtype: jnp.dtype = jnp.bfloat16,