Fix Ruff CI

src/ntops/kernels/__init__.py

@@ -1,9 +1,17 @@
 from ntops.kernels import (
     abs,
+    acosh,
+    adaptive_avg_pool2d,
+    adaptive_max_pool2d,
     add,
     addmm,
+    addmv,
     alpha_dropout,
+    argsort,
+    atan,
     avg_pool2d,
+    batch_norm,
+    bincount,
     bitwise_and,
     bitwise_not,
     bitwise_or,
@@ -18,6 +26,7 @@
     dropout,
     eq,
     exp,
+    fmax,
     ge,
     gelu,
     gt,
@@ -26,8 +35,18 @@
     isnan,
     layer_norm,
     le,
+    logsumexp,
+    lp_pool1d,
+    lp_pool2d,
+    lp_pool3d,
     lt,
+    max,
+    max_pool1d,
     max_pool2d,
+    max_pool3d,
+    maximum,
+    mean,
+    median,
     mm,
     msort,
     mul,
@@ -51,29 +70,10 @@
     sin,
     softmax,
     sort,
+    stack,
     sub,
     tanh,
     threshold,
-    max_pool1d,
-    max_pool3d,
-    stack,  
-    mean,
-    median,
-    maximum,
-    atan,
-    batch_norm,
-    bincount,
-    adaptive_max_pool2d,  
-    acosh,
-    adaptive_avg_pool2d,
-    addmv,
-    argsort,
-    fmax,
-    logsumexp,
-    lp_pool1d,
-    lp_pool2d,
-    lp_pool3d,
-    max,
 )
 
 __all__ = [

src/ntops/kernels/acosh.py

@@ -21,66 +21,64 @@ def _log(x, dtype):
 def application(input, output):
     """
     计算反双曲余弦函数 acosh(x) = ln(x + sqrt(x² - 1))
-    
+
     参数:
     input: 输入张量，形状为 (C // block_size, block_size)
     output: 输出张量，形状为 (C // block_size, block_size)
-    
+
     数值稳定性考虑:
     1. 当x接近1时，x² - 1接近0，使用(x-1)(x+1)形式避免精度损失
     2. 当x很大时，避免x²溢出，使用代数变换
     3. float16特殊处理，提升到float32计算
     """
     dtype = output.dtype.dtype
-    
+
     for i in range(input.shape[0]):
         # 获取当前块的数据
         input_block = ntl.cast(input[i], dtype)
-        
+
         # 数值稳定的acosh计算
         # acosh(x) = ln(x + sqrt(x² - 1))
-        
+
         # 处理x接近1的情况：x² - 1 = (x-1)(x+1)
         # 这样可以避免当x接近1时的精度损失
         x_minus_one = input_block - ntl.cast(1.0, dtype)
         x_plus_one = input_block + ntl.cast(1.0, dtype)
-        
+
         # 计算 sqrt(x² - 1) = sqrt((x-1)(x+1))
         sqrt_term = _sqrt(x_minus_one * x_plus_one, dtype)
-        
+
         # 计算 x + sqrt(x² - 1)
         # 当x很大时，这可能会导致数值问题，但acosh的定义域x≥1
         sum_term = input_block + sqrt_term
-        
+
         # 最终计算 ln(x + sqrt(x² - 1))
         result = _log(sum_term, dtype)
-        
+
         # 处理边界情况：当x < 1时，acosh未定义，返回NaN
         # 当x == 1时，acosh(1) = 0
         result = ntl.where(
             input_block < ntl.cast(1.0, dtype),
             ntl.cast(float("nan"), dtype),
             ntl.where(
-                input_block == ntl.cast(1.0, dtype),
-                ntl.cast(0.0, dtype),
-                result
-            )
+                input_block == ntl.cast(1.0, dtype), ntl.cast(0.0, dtype), result
+            ),
         )
-        
+
         # 将结果存入输出
         output[i] = result
 
 
 def premake(ndim, dim, dtype=None, block_size=None):
     """
     准备acosh内核
-    
+
     参数:
     ndim: 输入张量的维度
     dim: 要计算acosh的维度
     dtype: 数据类型
     block_size: 分块大小，用于优化内存访问
-    
+
     返回:
     arrangement_: 张量排列函数
     application: 计算函数
@@ -89,8 +87,8 @@ def premake(ndim, dim, dtype=None, block_size=None):
     arrangement_ = functools.partial(arrangement, dim=dim, block_size=block_size)
 
     tensors = (
-        Tensor(ndim, dtype=dtype),           # 输入张量
-        Tensor(ndim, dtype=dtype),           # 输出张量
+        Tensor(ndim, dtype=dtype),  # 输入张量
+        Tensor(ndim, dtype=dtype),  # 输出张量
     )
 
-    return arrangement_, application, tensors
+    return arrangement_, application, tensors

src/ntops/kernels/adaptive_avg_pool2d.py

@@ -1,32 +1,43 @@
 import functools
+
 import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
+
 def _cast_to_f32(x, dtype):
     """
     为了保证累加精度，如果是 float16 则转为 float32 计算
     """
     return ntl.cast(x, ntl.float32) if dtype == ntl.float16 else x
 
-def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size):
+
+def arrangement(
+    input,
+    output,
+    kernel_size_flatted,
+    kernel_size_h,
+    kernel_size_w,
+    stride_h,
+    stride_w,
+    block_size,
+):
     if block_size is None:
         block_size = ninetoothed.block_size()
 
-    # input: (N, C, H_in, W_in) 
+    # input: (N, C, H_in, W_in)
     # output: (N, C, H_out, W_out)
 
     # 1. 窗口切分
     input_arranged = input.tile(
-        (1, 1, kernel_size_h, kernel_size_w), 
-        (1, 1, stride_h, stride_w)
+        (1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w)
     )
     # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w)
 
     # 2. 展平与重排
     input_arranged = input_arranged.ravel()
     # => (N, C, H_out, W_out, 1, 1, k_h, k_w)
-    
+
     input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1)
     # => (N*C*H_out*W_out, k_h*k_w)
 
@@ -35,7 +46,7 @@ def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w
     nearest_pow2 = 1 << (kernel_size_h * kernel_size_w - 1).bit_length()
     input_arranged = input_arranged.tile((1, nearest_pow2))
     # => (..., 1), dtype=(1, nearest_pow2)
-    
+
     input_arranged.dtype = input_arranged.dtype.squeeze(0)
     input_arranged = input_arranged.tile((block_size, -1))
     input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1)
@@ -50,28 +61,32 @@ def arrangement(input, output, kernel_size_flatted, kernel_size_h, kernel_size_w
 
     return input_arranged, output_arranged, kernel_size_flatted
 
+
 def application(input, output, kernel_size_flatted):
     # input: (block_size, nearest_pow2)
     # output: (block_size, )
     # kernel_size_flatted: scalar tensor (k_h * k_w)
 
     dtype = input.dtype
-    
+
     # 转为高精度进行 Sum
     val = _cast_to_f32(input, dtype)
-    
+
     # 求和 (Axis 1 对应 nearest_pow2 维度)
     # 这里的 0 填充不会影响 Sum 结果
     acc = ntl.sum(val, axis=1)
-    
+
     # 求平均： Sum / Area
     # 注意：kernel_size_flatted 是实际的窗口大小，不是 nearest_pow2
     res = acc / ntl.cast(kernel_size_flatted, acc.dtype)
-    
+
     # 转回原类型
-    output = ntl.cast(res, dtype)
+    output = ntl.cast(res, dtype)  # noqa: F841
+
 
-def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None):
+def premake(
+    ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None
+):
     arrangement_ = functools.partial(
         arrangement,
         kernel_size_h=kernel_size_h,
@@ -83,9 +98,9 @@ def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=N
 
     tensors = (
         # input: 设置 other=0，保证 tile 补齐的值不影响 sum
-        Tensor(ndim, dtype=dtype, other=0),     
-        Tensor(ndim, dtype=dtype),              # output
-        Tensor(0, dtype=dtype),                 # kernel_size_flatted (scalar)
+        Tensor(ndim, dtype=dtype, other=0),
+        Tensor(ndim, dtype=dtype),  # output
+        Tensor(0, dtype=dtype),  # kernel_size_flatted (scalar)
     )
 
-    return arrangement_, application, tensors
+    return arrangement_, application, tensors

src/ntops/kernels/adaptive_max_pool2d.py

@@ -1,37 +1,40 @@
 import functools
+
 import ninetoothed
 import ninetoothed.language as ntl
 from ninetoothed import Tensor
 
-def arrangement(input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size):
+
+def arrangement(
+    input, output, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size
+):
     if block_size is None:
         block_size = ninetoothed.block_size()
 
-    # input: (N, C, H_in, W_in) 
+    # input: (N, C, H_in, W_in)
     # output: (N, C, H_out, W_out)
-    
+
     # 使用 tile 将输入切分为窗口
     # floor_mode=True 对应默认行为，对于 Adaptive Pool，我们通常通过计算好的 stride/kernel 确保覆盖
     input_arranged = input.tile(
-        (1, 1, kernel_size_h, kernel_size_w), 
-        (1, 1, stride_h, stride_w)
+        (1, 1, kernel_size_h, kernel_size_w), (1, 1, stride_h, stride_w)
     )
     # => (N, C, H_out, W_out), dtype=(1, 1, k_h, k_w)
-    
+
     input_arranged = input_arranged.ravel()
     # => (N, C, H_out, W_out, 1, 1, k_h, k_w)
-    
+
     input_arranged = input_arranged.flatten(end_dim=4).flatten(start_dim=1)
     # => (N*C*H_out*W_out, k_h*k_w)
 
     # 找到最近的 2 的倍数用于并行规约
     nearest_pow2 = 1 << (kernel_size_h * kernel_size_w - 1).bit_length()
     input_arranged = input_arranged.tile((1, nearest_pow2))
     # => (..., 1), dtype=(1, nearest_pow2)
-    
+
     input_arranged.dtype = input_arranged.dtype.squeeze(0)
     # => (..., 1), dtype=(nearest_pow2, )
-    
+
     input_arranged = input_arranged.tile((block_size, -1))
     # => (..., 1), dtype=(block_size, nearest_pow2)
     input_arranged.dtype = input_arranged.dtype.ravel().squeeze(1)
@@ -45,19 +48,23 @@ def arrangement(input, output, kernel_size_h, kernel_size_w, stride_h, stride_w,
 
     return input_arranged, output_arranged
 
+
 def application(input, output):
-    # input: (block_size, nearest_pow2) 
+    # input: (block_size, nearest_pow2)
     # output: (block_size, )
-    
+
     # 简单的 max reduction
     # 因为在 premake 中设置了 other=float("-inf")，padding 部分的值为负无穷，
     # 或者是 nearest_pow2 补齐产生的部分，通常默认为 0 或 padding 值，
     # 这里为了安全，可以显式处理 padding，但如果 arrange padding 正确，直接 max 即可。
     # 假设 DSL 的 tile 填充行为遵循 Tensor 的 other 属性。
-
-    output = ntl.max(input, axis=1)
 
-def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None):
+    output = ntl.max(input, axis=1)  # noqa: F841
+
+
+def premake(
+    ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=None, dtype=None
+):
     arrangement_ = functools.partial(
         arrangement,
         kernel_size_h=kernel_size_h,
@@ -69,8 +76,8 @@ def premake(ndim, kernel_size_h, kernel_size_w, stride_h, stride_w, block_size=N
 
     tensors = (
         # input: 设置 other 为负无穷，这样 tile 越界填充的值不会影响 max
-        Tensor(ndim, dtype=dtype, other=float("-inf")), 
-        Tensor(ndim, dtype=dtype),              # output
+        Tensor(ndim, dtype=dtype, other=float("-inf")),
+        Tensor(ndim, dtype=dtype),  # output
     )
 
-    return arrangement_, application, tensors
+    return arrangement_, application, tensors

src/ntops/kernels/addmv.py

@@ -23,7 +23,7 @@ def arrangement(
 
     # 关键：强制 block_size_n 为 1，因为这是 MV 乘法
     if block_size_n is None:
-        block_size_n = 1 
+        block_size_n = 1
 
     if block_size_k is None:
         block_size_k = mm.BLOCK_SIZE_K
@@ -68,7 +68,7 @@ def application(input, mat, vec, beta, alpha, output, input_precision):
     # 这里 output 是 (M, 1)，zeros 也创建 (M, 1)
     mm_output = ntl.zeros(output.shape, dtype=ntl.float32)
     mm.application(mat, vec, mm_output, input_precision)
-    
+
     # 逐元素操作，形状兼容 (M, 1)
     output = beta * input + alpha * mm_output
 
@@ -102,4 +102,4 @@ def premake(
         Tensor(0, dtype=dtype, constexpr=True, value=input_precision),
     )
 
-    return arrangement_, application, tensors
+    return arrangement_, application, tensors