[JAX] Fix GroupedScaledTensor creation with keyword arg (#2154)

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

[JAX] Fix GroupedScaledTensor creation with keyword arg (#2154)
Fix GroupedScaledTensor creation Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
5b3d65cc · Phuong Nguyen · GitHub · c47f329b · 5b3d65cc
Unverified Commit 5b3d65cc authored Sep 08, 2025 by Phuong Nguyen Committed by GitHub Sep 08, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

transformer_engine/jax/dense.py transformer_engine/jax/dense.py +2 -2

No files found.
--- a/transformer_engine/jax/dense.py
+++ b/transformer_engine/jax/dense.py
@@ -442,7 +442,7 @@ def _grouped_dense_fwd_rule(
            ctx_kernel = ScaledTensorFactory.create_1x(
                global_ctx_kernel_data.reshape(-1),
                ctx_kernel.scale_inv,
-                ctx_kernel.scaling_mode,
+                scaling_mode=ctx_kernel.scaling_mode,
                dq_dtype=ctx_kernel.dq_dtype,
                is_colwise=False,
                data_layout="N",
@@ -459,7 +459,7 @@ def _grouped_dense_fwd_rule(
                grouped_gemm_kernel = ScaledTensorFactory.create_1x(
                    grouped_gemm_kernel_data.reshape(-1),
                    ctx_kernel.scale_inv,
-                    ctx_kernel.scaling_mode,
+                    scaling_mode=ctx_kernel.scaling_mode,
                    dq_dtype=ctx_kernel.dq_dtype,
                    is_colwise=True,
                    data_layout="T",