checkstyle

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -1,7 +1,10 @@
 import torch
 import triton
 
-from liger_kernel.ops.cross_entropy import element_mul_kernel, liger_cross_entropy_kernel
+from liger_kernel.ops.cross_entropy import (
+    element_mul_kernel,
+    liger_cross_entropy_kernel,
+)
 
 # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
@@ -155,7 +158,7 @@ def fused_linear_cross_entropy_backward(
             V = grad_bias.shape[0]
             n_rows = V
 
-            element_mul[(n_rows,)](
+            element_mul_kernel[(n_rows,)](
                 grad_bias,
                 grad_bias.stride(-1),
                 grad_output,

src/liger_kernel/ops/layer_norm.py

@@ -130,7 +130,6 @@ def _layer_norm_backward_kernel(
     tl.store(DB_ptr + row_block_id * stride_db + cols, db_row.to(dtype), mask=mask)
 
 
-
 def layer_norm_forward(X, W, B, eps):
     shape = X.shape
     dim = shape[-1]

src/liger_kernel/ops/rms_norm.py

@@ -291,7 +291,9 @@ def forward(ctx, X, W, eps, offset=0.0, casting_mode="llama"):
         X: (B, T, H) or (BxT, H)
         W: (H,)
         """
-        Y, r, BLOCK_SIZE, num_warps, casting_mode = rms_norm_forward(X, W, eps, offset, casting_mode)
+        Y, r, BLOCK_SIZE, num_warps, casting_mode = rms_norm_forward(
+            X, W, eps, offset, casting_mode
+        )
         ctx.eps = eps
         ctx.offset = offset
         ctx.casting_mode = casting_mode