jax/ep: apply clang-format and silence pylint unused-arg in lowering

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

Author: phu0ngng

transformer_engine/jax/cpp_extensions/ep.py

@@ -514,7 +514,7 @@ def lowering(
         out_leading_shape,
         out_partition_spec,
     ):
-        del out_partition_spec
+        del out_leading_shape, out_partition_spec
         return ffi.ffi_lowering(EpCombinePrimitive.name)(
             ctx,
             handle_mem,
@@ -650,7 +650,7 @@ def lowering(
         out_leading_shape,
         out_partition_spec,
     ):
-        del out_partition_spec
+        del out_leading_shape, out_partition_spec
         return ffi.ffi_lowering(EpDispatchBwdPrimitive.name)(
             ctx,
             handle_mem,

transformer_engine/jax/csrc/extensions/ep.cpp

@@ -446,8 +446,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpDispatchBwdHandler, EpDispatchBwdFFI,
 // ── ep_combine_bwd ────────────────────────────────────────────────────────────
 
 Error_Type EpCombineBwdFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type handle_mem,
-                           Buffer_Type grad, Result_Type grad_expert_out,
-                           EpConfig config) {
+                           Buffer_Type grad, Result_Type grad_expert_out, EpConfig config) {
   (void)ep_state;
   auto grad_dims = grad.dimensions();
   NVTE_CHECK(grad_dims.size() >= 2,