照搬 JAX 的批维形状，静默算成别的批

主线位置：张量主线的第二步「计算」，批量线代和带批维的收缩。

版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native）

flowchart TB
  x["形状 4,2,3"] --> js["JAX svd"]
  x --> ts["tenferro svd"]
  js --> j4["4 个 2x3"]
  ts --> t3["3 个 4x2"]
  x --> tp["transpose 1,2,0"]
  tp --> ok["4 个 2x3"]

实验代码：JAX 脚本一个，tenferro 程序 p02_batch_axis（traced，CpuEngine::serial()，release）。两边用同一个公式生成数据：x[b, i, j] = cos(b + 2i + 5j)、y[b, i, j] = sin(b + 3i + 7j)。平台 M2 Max、macOS 27、rustc 1.99.0。

现象

JAX 的批量线代把批维放在最左边：

xs = np.array([[[np.cos(b + 2 * i + 5 * j) for j in range(3)] for i in range(2)] for b in range(4)])
s = jnp.linalg.svd(jnp.asarray(xs), compute_uv=False)

svd: x.shape (4, 2, 3) -> s.shape (4, 2)
s = [[1.5950948981274058, 0.8329548815374094], [1.5020504511629467, 0.8845522338360555], [1.295088064018098, 1.0259086766566576], [1.572363292005573, 0.8449968837773724]]

搬到 tenferro 时，数据按行主序导入（这一步是对的，见同一段数据 reshape 出来的矩阵不一样），形状照抄 [4, 2, 3]：

let x = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_row_major(vec![4, 2, 3], xs.clone())?)?;
let (_, s, _) = x.svd()?;

svd: s.shape [2, 3]  s = [1.757821914396673, 1.0080821312459942, 1.5543234985993488, 1.1400643839025426, 1.8613608089526557, 0.9520071838264343]

没有报错。tenferro 把前两维当矩阵、最后一维当批：3 个 4×2 矩阵，每个 2 个奇异值，没有一个数和 JAX 对得上。

带批维的 dot_general 也一样。JAX 的 (4,2,3) × (4,3,2)，批轴 0、收缩轴 2 和 1：

o = jax.lax.dot_general(xs, ys, (((2,), (1,)), ((0,), (0,))))

lax.dot_general (4,2,3) x (4,3,2), batch 0 -> shape (4, 2, 2)  o[1] = [[-0.2691029311765588, -0.4212242123100276], [-0.12675479794978153, 0.06496482486068857]]

tenferro 用同样的轴号：

let o = x.dot_general(&y, dg(&[2], &[1], &[0], &[0]))?;

output shape [2, 2, 4]
o[:, :, 1] rows [[-0.26910293117655876, -0.4212242123100276], [-0.1267547979497815, 0.06496482486068855]]
same buffer read as JAX's (4, 2, 2), [1] rows [[-0.1293951820770959, 0.4146305574433564], [-0.1267547979497815, -0.03860143268356753]]

数值是对的（第 1 批和 JAX 最大差 5.6e-17），但输出形状是 [2, 2, 4]：批轴在最后。下游代码如果还按 JAX 的 (4, 2, 2) 去取第 1 批，取到的是别的元素（第三行）。

出处

两处是同一个约定：tenferro 的批维在右边。

批量线代：输入是 [M, N, B...]，每个批次的矩阵在列主序下占一段连续内存。为什么这样定见 tenferro 卷的列主序与视图，和 JAX 的并排对照见布局与 reshape 的批维一节。
TracedTensor::dot_general 的文档写着「The output layout is [lhs free..., rhs free..., batch...]: batch axes come last」。lax.dot_general 是「批、左自由、右自由」。轴序对照见 einsum 的 dot_general 一节。

形状 [4, 2, 3] 本身就是合法的批量矩阵，tenferro 没有办法知道你想要的批维是哪一个。

修法

导入之后把批维挪到右边，算完再挪回来（只在交给 JAX 习惯的代码时才需要挪回）：

let xr = x.transpose(&[1, 2, 0])?;      // [4, 2, 3] -> [2, 3, 4]
let (_, s2, _) = xr.svd()?;             // s: [2, 4]
let o2 = o.transpose(&[2, 0, 1])?;      // dot_general 输出 [2, 2, 4] -> [4, 2, 2]

svd: s.shape [2, 4]
s[:, b] per batch = [[1.595094898127406, 0.8329548815374095], [1.5020504511629464, 0.8845522338360555], [1.2950880640180977, 1.0259086766566576], [1.572363292005573, 0.8449968837773721]]
...
o.transpose([2, 0, 1]) shape [4, 2, 2], [1] rows [[-0.26910293117655876, -0.4212242123100276], [-0.1267547979497815, 0.06496482486068855]]

奇异值和 JAX 最大差 2.2e-16（量级 1.6），两个不同的 SVD 实现舍入不同；dot_general 第 1 批最大差 5.6e-17。新写的 tenferro 代码最好从一开始就把批维放在右边，只在边界上转换一次。

怎么确认没中招

打印批量运算输出的形状：svd 的 s 应该是 [min(M, N), B]，dot_general 的输出最后一维应该是批长度。上面错的那次 s.shape 是 [2, 3]，批长度 4 根本没出现。
选一个批次单独算（不带批维的 2×3 矩阵），和批量结果的对应切片比较。
批长度和矩阵某一维相等时（例如 3 个 3×3 矩阵），形状检查查不出来；测试时让批长度和矩阵维度都不相同。

小结

JAX 批维在左（(B, M, N)），tenferro 批维在右（[M, N, B]），dot_general 的输出也是批轴最后。
照抄形状不会报错，只会换一种方式把数据切成矩阵，算出另一批结果。
导入后 transpose 把批维挪到右边；dot_general 的输出按「左自由、右自由、批」读。

照搬 JAX 的批维形状，静默算成别的批

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录