Basalt 3.11.0 — 按轴归约

Basalt 3.11.0

这一版补的是按轴归约：没有它，数值稳定的 softmax 写不出来。

• 坐标 — io.github.9178nekomc:basalt:3.11.0
• 在线文档 — https://9178nekomc.github.io/basalt/
• 运行要求 — Java 21，启动加 --add-modules jdk.incubator.vector
• 规模 — 243 个公开类型 / 1607 个公开成员，全部带逐参数说明

为什么

3.10.0 加了全量归约，但不能减最大值。这意味着 softmax 在输入偏大时直接溢出：
exp(1000) = inf，inf / inf = NaN。上一版的「手写 softmax」只在输入温和时成立。

新增

算子	反向规则
sum(dim) / sum(dim, keepdim)	梯度原样广播到该维每个位置
mean(dim) / mean(dim, keepdim)	梯度除以该维长度后广播
max(dim) / max(dim, keepdim)	只给到该组第一个最大值的位置，并列取最早

原生算子 31 → 34。keepdim 保留成长度 1，是为了能直接和原张量广播。

现在稳定的 softmax 是四行，且在 1000 量级输入上不溢出：

Tensor shifted = x.sub(x.max(1, true));
Tensor numer = shifted.exp();
Tensor denom = numer.sum(1, true);
Tensor probs = numer.div(denom);

一个形状决策

库里没有 0 维张量。rank-1 张量做 sum(0) 时，PyTorch 会给一个 0 维标量，
这里统一折成形状 [1]——和 sumAll 保持一致，也免得每个调用方都要处理空形状。

验证

• 28 个 Rust 单元测试（新增 3 个：前向值、梯度路由、非法形状）
• PrimitiveOpsTest 10 项，新增：
  • 三个归约的数值与 keepdim 形状
  • 越界维度抛异常
  • 三个归约各自的中心差分梯度检验
  • [1000, 1001, 1002] 的 softmax 必须有限且和为 1——不做 max 归约就会 NaN，这条是本次的判据
• ./verify.sh 全绿（9 组套件）
• 文档：19 页 / 243 类型 / 1607 个成员全部有说明 / 0 占位符 / 0 断链

还没做的

• 索引算子：gather / slice / cat — 手写注意力头、embedding 变体需要
• 掩码：where / maskedFill — 变长序列与 causal mask 需要
• 多轴归约：一次压掉多个维度，现在只能一维一维来
• 自定义算子注册：用 Java 写前向反向、不改 Rust。前置条件是延迟执行

产物

文件	说明
basalt-3.11.0.jar	标准 jar，含四平台原生库
basalt-3.11.0-all.jar	Fat jar，无外部依赖
basalt-3.11.0-sources.jar	源码
basalt-docs-3.11.0.zip	离线文档，19 页