Pypto Case Reduction Sum

单轴归约：Sum Reduction（3D）

最简单的模式——不需要 loop/view/assemble，kernel 只有 3 行。

def create_sum_reduction_kernel(in_shape, out_shape):
    @pypto.frontend.jit(runtime_options=..., debug_options=...)
    def kernel(
        x: pypto.Tensor(in_shape, pypto.DT_FP32),
    ) -> pypto.Tensor(out_shape, pypto.DT_FP32):
        output = pypto.tensor(list(out_shape), pypto.DT_FP32)
        pypto.set_vec_tile_shapes(1, 16, 256)
        output[:] = pypto.sum(x, dim=1, keepdim=True)
        return output
    return kernel

forward：保持原始维度，不降维。

def forward(self, x):
    assert x.dim() == 3
    assert tuple(x.shape) == (16, 256, 256)
    assert self.dim == 1
    x = x.contiguous()
    batch, _, dim2 = x.shape
    return create_sum_reduction_kernel(
        tuple(x.shape), (batch, 1, dim2)
    )(x)

模式要点

• 保持输入原始维度，不 reshape → tile 参数个数 = input rank
• kernel 极简：set_tile + sum + return，无需 loop/view/assemble
• pypto.amin / pypto.amax 同理，只换 API
• mean = sum * (1.0 / size)（无内建 mean API）
• 对 (16, 256, 256), dim=1 的 3D 单轴归约，默认首选从 (1, 16, 256) 起步，再按 32/64 对照实测。
• 这里不要套用 loop 的“中段先试”习惯到 tile；该固定形状直接以 (1, 16, 256) 作为默认实现。
• 对该固定形状，(1, 32, 256) 与 (1, 64, 256) 不作为默认模板，仅作为对照候选。

强约束（reduction_over_a_dimension 系列）

• 本系列题目的 get_init_inputs() 返回值是本次固定参数（例如 dim=1）。
• 注释里的 Example, change to desired dimension 是题库说明，不是当前实现目标。
• 生成代码时：
  • 保留 ModelNew.__init__(dim) 签名；
  • forward 中 assert self.dim == <固定值>；
  • kernel 使用固定常量 dim=<固定值>，不要写 if dim == ... 分支；
  • 固定 dim 场景下，create_*_kernel 不再接收 dim 运行时参数。

反例（不要这样写）：

def create_xxx_kernel(in_shape, out_shape, dim):
    ...
    output[:] = pypto.amin(x, dim=dim, keepdim=True)

正例（固定 dim 写死）：

FIXED_DIM = 1
def create_xxx_kernel(in_shape, out_shape):
    ...
    output[:] = pypto.amin(x, dim=FIXED_DIM, keepdim=True)