fix lint errors

Signed-off-by: Keshav <[email protected]>

transformer_engine/jax/cpp_extensions.py

@@ -642,7 +642,8 @@ def infer_sharding_from_operands(zero_centered_gamma, epsilon, mesh, arg_infos,
         g_b_spec = get_padded_spec(arg_infos[4])
         if g_b_spec[-1] is not None:
             warnings.warn(
-                f"{LayerNormBwdPrimitive.name} does not support sharding of gradients of gamma and beta of Layernorm" \
+                f"{LayerNormBwdPrimitive.name} does not support sharding of gradients " \
+                f"of gamma and beta of Layernorm " \
                 f"Enforcing no sharding of parameters hidden dim! " \
             )
 
@@ -663,10 +664,11 @@ def partition(zero_centered_gamma, epsilon, mesh, arg_infos, result_infos):
         g_b_spec = get_padded_spec(arg_infos[4])
         if g_b_spec[-1] is not None:
             warnings.warn(
-                f"{LayerNormBwdPrimitive.name} does not support sharding of gradients of gamma and beta of Layernorm" \
+                f"{LayerNormBwdPrimitive.name} does not support sharding of gradients " \
+                f"of gamma and beta of Layernorm " \
                 f"Enforcing no sharding of parameters hidden dim! " \
             )
-        
+
         dx_sharding = NamedSharding(mesh, PartitionSpec(*x_spec[:-1], None))
         dgamma_sharding = dbeta_sharding = NamedSharding(mesh, PartitionSpec(None))
         out_shardings = dx_sharding, dgamma_sharding, dbeta_sharding
@@ -857,7 +859,7 @@ def partition(epsilon, mesh, arg_infos, result_infos):
                 f"{RmsNormFwdPrimitive.name} does not support sharding of parameter gamma " \
                 f"Enforcing no sharding of parameters hidden dim! " \
             )
-        
+
         x_sharding = NamedSharding(mesh, PartitionSpec(*x_spec[:-1], None))
         g_sharding = NamedSharding(mesh, PartitionSpec(None))
         out_sharding = x_sharding