第 3 课：张量运算扩展 —— 广播、归约、softmax、批量矩阵乘法

代码位置：src/tensor.rs（本章新增：broadcast_shapes / broadcast_map / binary / sum / sum_last_dim / softmax_last_dim / permute / gather_rows / 3D matmul） 演示入口：src/main.rs

1. 本课要搞懂的问题

1. 两个形状不同的张量怎么直接相加？"广播"到底是什么规则？
2. "求和"这个简单的操作，反向传播时梯度应该怎么分？
3. softmax 为什么要"先减最大值"？不减会出什么问题？
4. 形状重排（permute）、按行取数（gather_rows）的梯度怎么传回去？
5. 一批矩阵同时做乘法（3D 批量 matmul）是怎么实现的？

2. 广播 Broadcasting：形状不同也能做运算

2.1 为什么需要广播

上一课的 add 要求两个张量形状完全相同。但神经网络里处处是"小形状 + 大形状"的运算：

• 给一批数据加偏置：[B, D] + [D]
• 给一批样本各自乘一个标量：[B, D] * []
• 线性层 y = xW + b：[B, D] @ [D, out] + [out]

如果每次都要求形状一致，我们就得手动把 [D] 复制成 [B, D] 再算。广播（broadcasting） 就是自动完成这件事：只要形状"兼容"，Rust 帮你把小张量"展开"到大张量的形状再逐元素运算。

2.2 numpy 广播的三条规则

我们的实现完全遵循 numpy 的广播规则：

规则	说明	例子
① 右对齐	从最后一个维度开始，逐对比较两个形状的维度	[2, 3] 与 [3] → 对齐成 [2,3] vs [1,3]
② 相等或为 1	每一对维度要么相等，要么其中一个为 1（1 可以扩展成另一个）	2 vs 1 → 结果 2
③ 都不满足则失败	两个都大于 1 且不相等 → 无法广播，报错	[2, 3] 与 [2, 4] → None

几个直观例子（1 表示"这一维可以复制"）：

a 形状	b 形状	逐对比较	广播结果
[2, 3]	[3]	2 vs 1（b 补 1），3 vs 3	[2, 3]
[2, 3]	[2, 1]	2 vs 2，3 vs 1	[2, 3]
[3]	[1]	3 vs 1	[3]
[2, 3]	[2, 4]	3 vs 4 都不为 1	❌ 失败

2.3 broadcast_shapes：把规则变成代码

/// 计算两个形状广播后的形状（numpy 广播规则，从右向左对齐）：
/// - 维度相等或其中一个为 1 即可广播
fn broadcast_shapes(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
    let n = a.len().max(b.len());
    let mut out = vec![1usize; n];
    for i in 0..n {
        // 从右往左对齐：i 从 0 数到 n-1 时，i+n-1 正好指向"倒数第 n-i 个维度"
        // 如果 i + len 超出本形状的维度范围，就当作"前面补了 1"
        let da = if i + a.len() >= n { a[i + a.len() - n] } else { 1 };
        let db = if i + b.len() >= n { b[i + b.len() - n] } else { 1 };
        if da == db {
            out[i] = da;
        } else if da == 1 {
            out[i] = db;
        } else if db == 1 {
            out[i] = da;
        } else {
            return None;   // 规则 ③：无法广播
        }
    }
    Some(out)
}

要点：

• i + a.len() >= n 判断"对齐后当前位是否越过了 a 的最左端"，越过了说明 a 在这一维没有对应维度，视作补 1（规则 ① 的"右对齐"实现）。
• 三条分支对应规则 ②：相等取它、一方为 1 取另一方、都不满足返回 None。
• 返回 Option 而不是直接 panic，把"能不能广播"的判定权交给调用方。

验证一下 [2,3] 与 [3]（n = 2）：

i	da	db	判定	out
0	a[0] = 2（0+2>=2）	b 越界 → 1	db == 1	2
1	a[1] = 3	b[0] = 3（1+1>=2）	相等	3

结果 [2, 3] ✓

2.4 broadcast_map：数据怎么搬

光知道"结果形状"还不够——结果每个位置的值要从源数据的哪里取。broadcast_map 就是干这个的：给定源形状 src 和目标形状 target，返回一张映射表 map，map[t] = 目标第 t 个元素对应的源展平下标。

/// 计算源形状 src 广播到目标形状 target 时，每个目标元素对应的源展平下标。
/// src 的维度必须 <= target 且右对齐；src 中大小为 1 的维度索引固定为 0。
fn broadcast_map(target: &[usize], src: &[usize]) -> Vec<usize> {
    let offset = target.len() - src.len();       // 右对齐：源前面补 offset 个"伪维度"
    let total: usize = target.iter().product();
    let mut map = vec![0usize; total];
    for flat in 0..total {
        // 第一步：把目标的一维下标反解成多维下标（行优先）
        let mut r = flat;
        let mut t_idx = vec![0usize; target.len()];
        for d in (0..target.len()).rev() {
            t_idx[d] = r % target[d];
            r /= target[d];
        }
        // 第二步：多维下标映射回源的多维下标，再展平成源的一维下标
        let mut s_flat = 0usize;
        for d in 0..src.len() {
            let td = t_idx[d + offset];          // 源的第 d 维对应目标的第 d+offset 维
            let sd = if src[d] == 1 { 0 } else { td };  // 源大小为 1 的维度，索引固定为 0
            s_flat = s_flat * src[d] + sd;       // 行优先展平
        }
        map[flat] = s_flat;
    }
    map
}

反解多维下标：行优先存储下，一维下标 flat 的分解是"从最后一个维度开始取余、整除"：

t_idx[d] = r % target[d];  r /= target[d];    // d 从大到小

比如目标形状 [2, 3]，flat = 4 → t_idx[1] = 4%3 = 1，r = 1 → t_idx[0] = 1，即位置 (1, 1)。

映射到源：src 的第 d 维对应目标第 d + offset 维；若源该维大小为 1，则无论目标索引多大，都取下标 0（这就是"复制"）。

例：src = [3] 广播到 target = [2, 3]（相当于每行复制同一个向量）：

目标位置 (i, j)	flat	映射到源下标	说明
(0, 0)	0	0	取 b[0]
(0, 1)	1	1	取 b[1]
(0, 2)	2	2	取 b[2]
(1, 0)	3	0	第二行还是取 b[0]
(1, 1)	4	1	取 b[1]
(1, 2)	5	2	取 b[2]

2.5 binary：把广播装进"逐元素运算"

broadcast_plan 先判断两个张量是否同形状：同形状就直接逐位对齐（无需映射），不同形状才计算两个 broadcast_map：

fn broadcast_plan(&self, other: &Tensor) -> (Vec<usize>, Option<Vec<usize>>, Option<Vec<usize>>) {
    if self.shape == other.shape {
        (self.shape.clone(), None, None)          // 同形状：没有 map
    } else {
        let target = broadcast_shapes(&self.shape, &other.shape)
            .unwrap_or_else(|| panic!("形状无法广播：{:?} vs {:?}", self.shape, other.shape));
        let map_a = if self.shape == target { None } else { Some(broadcast_map(&target, &self.shape)) };
        let map_b = if other.shape == target { None } else { Some(broadcast_map(&target, &other.shape)) };
        (target, map_a, map_b)
    }
}

add / sub / mul / div 四个运算全部复用同一个通用函数 binary：

fn binary(
    &self,
    other: &Tensor,
    fwd: impl Fn(f32, f32) -> f32 + 'static,          // 前向：c = f(a, b)
    back: impl Fn(f32, f32) -> (f32, f32) + 'static,  // 反向：返回 (∂c/∂a, ∂c/∂b)
) -> Tensor {
    let (target_shape, map_a, map_b) = self.broadcast_plan(other);
    let sa = self.data.borrow();
    let sb = other.data.borrow();
    let total: usize = target_shape.iter().product();
    let mut out_data = vec![0.0f32; total];
    for t in 0..total {
        let ia = match &map_a { Some(m) => m[t], None => t };   // 没有 map 就是一一对应
        let ib = match &map_b { Some(m) => m[t], None => t };
        out_data[t] = fwd(sa[ia], sb[ib]);                     // 按映射取数做运算
    }
    drop(sa);
    drop(sb);
    // …… 构造结果节点，注册反向闭包（见下）……
}

反向传播：广播的梯度遵循"谁被复制了，梯度就累加回谁"。因为一个源元素可能对应多个输出位置，∂loss/∂x = Σ_t (∂loss/∂out_t)·(∂out_t/∂x)，所以必须累加（+=）：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let (ga, gb) = if same_shape && Rc::ptr_eq(&sg, &og) {
        // 同一张量参与运算（如 x*x）：两条路径的梯度合并（第 2 课讲过）
        // ……（合并累加逻辑，返回 (true, true)）
        (true, true)
    } else {
        let mut sgm = sg.borrow_mut();
        let mut ogm = og.borrow_mut();
        for t in 0..g.len() {
            let ia = match &map_a_c { Some(m) => m[t], None => t };
            let ib = match &map_b_c { Some(m) => m[t], None => t };
            let (da, db) = back(sd_b[ia], od_b[ib]);
            sgm[ia] += g[t] * da;      // 梯度累加到"源位置"，而不是输出位置
            ogm[ib] += g[t] * db;
        }
        (false, false)
    };
    let _ = (ga, gb);
}));

用测试验证（test_broadcast_add）：

// [2,3] + [3] 广播：b 的每个元素被两行"共享"
let a = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]);
let b = Tensor::param(vec![10.0, 20.0, 30.0], vec![3]);
let loss = a.add(&b).sum();
loss.backward();
// b 的每个元素被广播两次，梯度应为 2
assert_eq!(b.grad(), vec![2.0, 2.0, 2.0]);
assert_eq!(a.grad(), vec![1.0; 6]);

b[0] 参与了输出的两个位置 (0,0) 和 (1,0)，两个位置的梯度各 1，累加起来就是 2。✓

3. 归约：sum / sum_last_dim / mean

"归约"（reduction）就是把一堆数合并成一个数。神经网络里损失函数 loss 几乎总是某种归约的结果（loss.backward() 要求标量，所以归约是自动微分的必经之路）。

3.1 sum：全部求和成标量

/// 求和成标量，梯度均匀传给每个元素
pub fn sum(&self) -> Tensor {
    let total = self.data.borrow().iter().sum();
    let mut result = Tensor::new(vec![total], vec![], self.requires_grad);
    if self.requires_grad {
        // ……
        result.backward = Some(Rc::new(move || {
            let g = rg.borrow()[0];
            let mut sgm = sg.borrow_mut();
            for v in sgm.iter_mut() {
                *v += g;          // 每个元素的梯度都是 g（因为 ∂sum/∂x_i = 1）
            }
        }));
    }
    result
}

反向逻辑一句话：求和把每个元素的梯度都变成同一个 g。因为 sum(x) = x₀ + x₁ + …，对每个 x_i 求导都是 1，所以 ∂loss/∂x_i = g·1。

3.2 sum_last_dim：沿最后一维求和，保持维度

scores [B, D] 想按行求和得到 [B, 1]（注意最后一维保留为 1，而不是塌缩成 [B]），就用它：

/// 沿最后一维求和，**保持维度**：[..., D] -> [..., 1]
/// 反向：梯度广播回最后一维
pub fn sum_last_dim(&self) -> Tensor {
    assert!(self.rank() >= 1, "sum_last_dim 需要至少 1 维");
    let (pre, d) = (self.numel() / self.shape[self.rank() - 1], self.shape[self.rank() - 1]);
    // pre = 最后一维之前的所有元素数，d = 最后一维大小
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; pre];
    for p in 0..pre {
        let mut s = 0.0;
        for j in 0..d {
            s += sd[p * d + j];   // 每 d 个连续元素求和
        }
        out_data[p] = s;
    }
    drop(sd);
    let mut new_shape = self.shape.clone();
    *new_shape.last_mut().unwrap() = 1;   // 最后一维改成 1，保持维度
    // ……
    result.backward = Some(Rc::new(move || {
        let g = rg.borrow();
        let mut sgm = sg.borrow_mut();
        for p in 0..pre {
            for j in 0..d {
                sgm[p * d + j] += g[p];   // 反向：把 g[p] 广播回这一行的 d 个位置
            }
        }
    }));
}

把 [2, 3] 变成 [2, 1]，前向和反向正好是"一对"：

前向： [1 2 3]   sum -> [6]
       [4 5 6]          [15]

反向： g = [2, 3]  broadcast -> [2 2 2]
                               [3 3 3]

3.3 mean：用 sum + mul_scalar 组合出来

我们的库没有单独的 mean 算子——均值就是"求和再除以元素数"，用现成的算子组合即可。看 loss.rs 里 MSE 的实现：

pub fn mse_loss(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub(target)
        .pow(2.0)
        .sum()                                  // Σ (pred - target)²
        .mul_scalar(1.0 / pred.numel() as f32)  // ÷ N，就是 mean
}

这体现了"先保证运算集合最小、用组合满足需求"的设计哲学：少一个算子 = 少一段前向 + 反向代码 = 少一类 bug。

三种归约对比：

运算	输入 → 输出	反向梯度
sum()	[D1, D2, …] → [] 标量	每个元素都 += g
sum_last_dim()	[…, D] → […, 1] 保持维度	把 g[p] 广播回第 p 行的 d 个位置
mean（组合实现）	同上	同 sum，再乘 1/N

4. softmax_last_dim：数值稳定的 softmax

4.1 公式与直觉

softmax 把一行任意实数分数变成概率分布（每个数在 0~1 之间，和为 1），是分类问题（以及 Transformer 注意力）的核心：

softmax(x)_i = e^{x_i} / Σ_j e^{x_j}      （对每一行独立做）

x	e^x	归一化后
1	2.718	0.0900
2	7.389	0.2447
3	20.086	0.6652

这就是测试 test_softmax 验证的结果：softmax([1,2,3]) = [0.0900, 0.2447, 0.6652]。

4.2 数值爆炸：不先减 max 会怎样

e^x 增长极快：e^88.7 就已经超过 f32 能表示的最大值（约 3.4×10³⁸），再大就是 inf。而神经网络的 logits 动辄上百，直接 exp 会溢出，除以 inf 得到 NaN，整个训练直接崩掉。

4.3 "减最大值"技巧

数学上 softmax 对"整体平移"不敏感，因为分子分母可以同乘一个常数：

softmax(x)_i = e^{x_i} / Σ_j e^{x_j}
             = e^{-m}·e^{x_i} / (e^{-m}·Σ_j e^{x_j})
             = e^{x_i - m} / Σ_j e^{x_j - m}         其中 m = max_j x_j

取 m = max(x) 后，x_i - m ≤ 0，e^{x_i-m} ∈ (0, 1]，永远不会溢出。代码里就是三步：找每行 max → exp(x - max) → 除以行和：

pub fn softmax_last_dim(&self) -> Tensor {
    assert!(self.rank() >= 1, "softmax_last_dim 需要至少 1 维");
    let (rows, d) = (self.numel() / self.shape[self.rank() - 1], self.shape[self.rank() - 1]);
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; rows * d];
    // 先存 softmax 结果（反向需要）
    for r in 0..rows {
        let mut maxv = f32::NEG_INFINITY;
        for j in 0..d {
            maxv = maxv.max(sd[r * d + j]);            // ① 找每行最大值
        }
        let mut sum = 0.0;
        for j in 0..d {
            out_data[r * d + j] = (sd[r * d + j] - maxv).exp();  // ② 减 max 再 exp
            sum += out_data[r * d + j];
        }
        for j in 0..d {
            out_data[r * d + j] /= sum;                // ③ 归一化
        }
    }
    drop(sd);
    // ……（反向见下）……
}

注：f32::NEG_INFINITY 作为"比任何数都小"的初始值，保证第一个元素就能更新 max。

4.4 softmax 的反向传播

softmax 的输出互相耦合（每个 s_i 都依赖所有 x_j），梯度公式（s 是输出，g 是输出梯度）：

∂loss/∂x_i = s_i · (g_i - Σ_j g_j·s_j)

记号	含义
s_i	softmax 输出的第 i 个分量
g_i	上一层传下来的梯度（对输出的梯度）
dot = Σ_j g_j·s_j	梯度与输出的点积（对每一行只算一次）

代码和公式一一对应：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for r in 0..rows {
        // dot = Σ_j g_j * s_j
        let mut dot = 0.0;
        for j in 0..d {
            dot += g[r * d + j] * out_data[r * d + j];
        }
        for i in 0..d {
            sgm[r * d + i] += out_data[r * d + i] * (g[r * d + i] - dot);  // s_i(g_i - dot)
        }
    }
}));

推导速览（供参考）：∂s_i/∂x_j = s_i(δ_ij - s_j)（δ 是克罗内克记号），再用链式法则 ∂loss/∂x_i = Σ_j g_j·∂s_j/∂x_i = Σ_j g_j·s_j(δ_ji - s_i) = g_i·s_i - s_i·Σ_j g_j·s_j。

5. permute：任意维重排

5.1 和 reshape 的区别

• reshape：只改解释方式，不改数据顺序。[1,2,3,4,5,6] 排成 2×3 还是 3×2，内存里还是那 6 个数。
• permute：真的把数据搬家。把某个维度整体挪到另一个位置（transpose 就是 2 维的 permute([1,0])）。

permute(&[0, 2, 1]) 的意思是"新张量的第 0 维 = 旧张量的第 0 维，新第 1 维 = 旧第 2 维，新第 2 维 = 旧第 1 维"。

看测试 test_permute：[2, 3] 用 permute([1, 0]) 变成 [3, 2]，数据从 [1,2,3,4,5,6] 变成 [1,4,2,5,3,6]：

原始 [2,3]:          permute([1,0]) 后 [3,2]:
[1 2 3]              [1 4]
[4 5 6]              [2 5]
                     [3 6]

5.2 用"逆映射"实现

要算 out[of] = src[?]，思路是：先反解 of 的多维坐标 out_idx，再问"源的第 d 维应该取输出的哪一维"——答案就是 inv[d]，其中 inv 是 perm 的逆映射：

// 反解 permute 的逆映射：inv[perm[i]] = i
let mut inv = vec![0usize; self.rank()];
for (i, &d) in dims.iter().enumerate() {
    inv[d] = i;
}
// 前向：out_flat -> src_flat
let mut map = vec![0usize; total];
for out_flat in 0..total {
    // ① 反解输出多维下标（行优先）
    let mut r = out_flat;
    let mut out_idx = vec![0usize; self.rank()];
    for d in (0..self.rank()).rev() {
        out_idx[d] = r % new_shape[d];
        r /= new_shape[d];
    }
    // ② 源的第 d 维来自输出的第 inv[d] 维
    let mut src_flat = 0usize;
    for d in 0..self.rank() {
        let sd = out_idx[inv[d]];
        src_flat = src_flat * self.shape[d] + sd;
    }
    map[out_flat] = src_flat;
}
// ③ 按映射搬运数据
let mut out_data = vec![0.0f32; total];
for (of, &sf) in map.iter().enumerate() {
    out_data[of] = sd[sf];
}

反向：和广播类似——permute 也是"每个源元素只去一个地方"，但方向反过来，梯度要"放回原位"。用同一张 map，把输出的梯度 g[of] 累加到 sgm[sf]：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for (of, &sf) in map.iter().enumerate() {
        sgm[sf] += g[of];        // 梯度"原路返回"
    }
}));

为什么是 +=？因为 permute 之后可能还有别的运算把梯度送到同一个源位置，或者源张量被多处共享，梯度必须累加（全项目统一约定）。

6. gather_rows：按行取数（Embedding 的基石）

6.1 语义

词嵌入（Embedding）的本质就是"查表"：一张 [V, D] 的表格（V 个词的向量），给定一堆词的下标，把对应行拿出来：

table [V, D] + indices [N]  ->  out [N, D]
out[i, :] = table[indices[i], :]

前向代码：

pub fn gather_rows(&self, indices: &[usize]) -> Tensor {
    assert_eq!(self.rank(), 2, "gather_rows 的 table 必须为 2 维");
    let (v, d) = (self.shape[0], self.shape[1]);
    let n = indices.len();
    let sd = self.data.borrow();
    let mut out_data = vec![0.0f32; n * d];
    for (i, &idx) in indices.iter().enumerate() {
        assert!(idx < v, "gather 索引越界：{} >= {}", idx, v);
        for j in 0..d {
            out_data[i * d + j] = sd[idx * d + j];   // 复制第 idx 行
        }
    }
    // ……
}

测试 test_gather_rows：table[3,2] 取第 0、2 行 → [1,2,5,6] ✓

6.2 反向：scatter-add（散射累加）

一行可能被取多次（两个位置都引用同一行），所以反向时梯度累加回表格对应行：

result.backward = Some(Rc::new(move || {
    let g = rg.borrow();
    let mut sgm = sg.borrow_mut();
    for i in 0..idx_vec.len() {
        let row = idx_vec[i];                 // 输出第 i 行来自表格第 row 行
        for j in 0..d2 {
            sgm[row * d2 + j] += g[i * d2 + j];   // 累加，不覆盖！
        }
    }
}));

测试验证：

let t = Tensor::param(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
let g = t.gather_rows(&[0, 2]);
let loss = g.sum();
loss.backward();
// 第 0 行梯度 1，第 2 行梯度 1，第 1 行梯度 0
assert_eq!(t.grad(), vec![1.0, 1.0, 0.0, 0.0, 1.0, 1.0]);

表格行	被引用次数	梯度
0	1 次（indices[0]）	1
1	0 次	0
2	1 次（indices[1]）	1

如果 indices = [0, 0]，第 0 行的梯度就是 2——这就是"scatter add"里 add 的含义。

7. 3D 批量 matmul

7.1 从 2D 到 3D

第 1 课学过 2D 矩阵乘法 C[m,n] = A[m,k] @ B[k,n]。真实模型里我们通常同时处理一批样本，就变成了 3D 批量版本：前两个维度完全一样（都是 B），第三个、第四个维度才是矩阵的 m、k、n：

2D:   C[m,n]     = A[m,k]     @ B[k,n]
3D:   C[B,m,n]   = A[B,m,k]   @ B[B,k,n]     每个 batch 独立做一次 2D 乘法

批量维度 B 必须一致，内维 k 也必须一致，输出是 [B, m, n]。

7.2 前向代码

就是"对每个 batch 跑一遍第 1 课的三重循环"：

pub fn matmul(&self, other: &Tensor) -> Tensor {
    assert!(
        (self.rank() == 2 && other.rank() == 2) || (self.rank() == 3 && other.rank() == 3),
        "matmul 只支持 2D 或 3D（批量），当前 {}-D x {}-D", self.rank(), other.rank()
    );
    if self.rank() == 2 {
        return self.matmul_2d(other);        // 2D 走老路
    }
    // 3D 批量
    assert_eq!(self.shape[0], other.shape[0], "批量维度必须一致");
    let (b, m, k1) = (self.shape[0], self.shape[1], self.shape[2]);
    let (_, k2, n) = (other.shape[0], other.shape[1], other.shape[2]);
    assert_eq!(k1, k2, "矩阵乘法维度不匹配");

    let sd = self.data.borrow();
    let od = other.data.borrow();
    let mut out_data = vec![0.0f32; b * m * n];
    for bi in 0..b {                      // 多了一层 batch 循环
        for i in 0..m {
            for j in 0..n {
                let mut s = 0.0;
                for k in 0..k1 {
                    s += sd[(bi * m + i) * k1 + k] * od[(bi * k1 + k) * n + j];
                }
                out_data[(bi * m + i) * n + j] = s;
            }
        }
    }
    // ……
}

下标规律：每个 batch 内的下标公式和 2D 完全一样，只是多加了 bi * m * k1（A）和 bi * k1 * n（B）的批内偏移。看测试 test_matmul_3d：

A [2,1,2] = [[1,2],[3,4]]      B [2,2,1] = [[5],[6]] , [[7],[8]]
C [2,1,1] = [[1·5+2·6=17], [3·7+4·8=53]]   →  [17, 53] ✓

7.3 反向公式（对每个 batch 独立成立）

第 2 课的反向公式在批量下逐 batch 套用：

输入	梯度公式	说明
A	∂A = g @ Bᵀ	输出梯度乘 B 的转置
B	∂B = Aᵀ @ g	A 的转置乘输出梯度

代码里就是三重循环写成 += 累加（sgm[(bi*m+i)*k1+k] += Σ_j g[...]·od[...]），并且保留了 Rc::ptr_eq 判断（x @ x 这种自乘场景两条路径合并）。注意维度细节：∂A 是 [B,m,k]（和 A 同形状），∂B 是 [B,k,n]，所以两个循环的累加下标不同：

// ∂A：sgm 用 (bi, i, k) 定位，对 j 求和
sgm[(bi * m + i) * k1 + k] += s;      // s = Σ_j g[bi,i,j] * B[bi,k,j]
// ∂B：ogm 用 (bi, k, j) 定位，对 i 求和
ogm[(bi * k1 + k) * n + j] += s;      // s = Σ_i A[bi,i,k] * g[bi,i,j]

8. 运行与测试

cargo test     # 全部测试通过（含新增：广播、3D matmul、softmax、permute、gather_rows）
cargo run      # 运行演示

重点测试与验证内容：

测试	验证什么
test_broadcast_add	[2,3] + [3] 前向正确，且 b 的梯度是 2（被广播两次）
test_matmul_3d	[2,1,2] @ [2,2,1] = [2,1,1]，数据 [17, 53]
test_softmax	输出和为 1，与手算值 [0.0900, 0.2447, 0.6652] 一致
test_permute	[2,3] permute 成 [3,2]，数据 [1,4,2,5,3,6]
test_gather_rows	查表结果正确，且只有被引用行有梯度

9. 动手练习

1. 手推一遍 [2,3] + [1,3] 的广播：画两个张量，标出 b 的每个元素被复制到哪些位置；再画反向时梯度如何累加回 b。
2. 用 sum_last_dim 和 mul_scalar 组合实现"按行求均值"（[B, D] -> [B, 1]），并推导梯度。
3. 设 d=2，手推 softmax 反向公式：s = softmax([x0, x1])，写出 ∂loss/∂x0 和 ∂loss/∂x1 的具体表达式。
4. 思考：permute 之后做 sum()，源张量的梯度是什么？为什么说 permute 只改变数据的"排列"而不改变"信息"？
5. 挑战题：gather_rows 若允许同一行被取两次（如 indices=[0,0]），反向时该行梯度是多少？这提示了 Embedding 的梯度为什么必须用 +=。

10. 本课总结

• 广播：右对齐、维度相等或一方为 1、否则失败；broadcast_shapes 定形状，broadcast_map 定搬数路线
• 广播的反向 = 梯度累加回源位置（谁被复制，谁收多份梯度）
• 归约：sum 变标量、sum_last_dim 保持维度；mean 用 sum + 除法组合
• softmax：减 max 保证数值稳定（数学等价、永不溢出），反向公式 s_i(g_i - dot)
• permute / gather_rows：都是"搬家"运算，反向用同一张映射表把梯度放回原位（累加）
• 3D matmul：批量维度不变，内层逐 batch 做 2D matmul，反向逐 batch 套 g@Bᵀ / Aᵀ@g
• 下一课：把项目拆成独立模块，并深入理解支撑这一切的 Rc<RefCell> 架构