同一段数据 reshape 出来的矩阵不一样
主线位置:张量主线的第一步「建数组」,先建一维数据再
reshape成矩阵的那一行。
版本:JAX 0.12.0.1(CPU,x64)、tenferro-rs 0.8.0(CPU,native)
flowchart TB d["0..5 一维数据"] --> jr["JAX reshape"] d --> tr["tenferro reshape"] jr --> rowfill["按行填"] tr --> colfill["按列填"] rowfill --> r1["210, 543"] colfill --> r2["420, 531"] tr --> fix["倒轴再倒回"] fix --> r1
实验代码:JAX 脚本一个,tenferro 程序 p01_reshape_order(traced,CpuEngine::serial(),release)。平台 M2 Max、macOS 27、rustc 1.99.0。
现象
JAX 里先建一维数组再 reshape 是最常见的建矩阵写法:
a = jnp.arange(6.0).reshape(2, 3)
v = jnp.array([1.0, 10.0, 100.0])
print(a.tolist(), (a @ v).tolist())
arange(6).reshape(2, 3) rows [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] a @ v = [210.0, 543.0]
照搬到 tenferro,一维张量的数据和 JAX 完全一样(一维没有行列之分),差别出在 reshape:
let x = TracedTensor::from_vec_col_major(vec![6], (0..6).map(|i| i as f64).collect())?;
let a = x.reshape([2, 3])?;
let av = a.dot_general(&v, dg(&[1], &[0], &[], &[]))?; // a @ v
rows [[0.0, 2.0, 4.0], [1.0, 3.0, 5.0]] a @ v = [420.0, 531.0]
(dg 是实验里填 DotGeneralConfig 的小函数,rows 按 d[i + m·j] 把列主序缓冲区按行打印。)形状对,dot_general 成功,结果是另一个矩阵乘出来的。三维也一样,JAX 的 arange(8).reshape(2,2,2).reshape(4,2) 是 [[0,1],[2,3],[4,5],[6,7]],tenferro 两次 reshape 得到 [[0,4],[1,5],[2,6],[3,7]]。
出处
JAX 的 reshape 默认 order="C":最后一维变得最快。tenferro 的张量只有列主序,TracedTensor::reshape 按第一维最快的顺序重新解释元素,相当于 jnp.reshape(.., order="F"),没有 order 参数。两边的对照和「一段程序里并轴、拆轴用同一种约定时结果与约定无关」这条规律,见布局与 reshape。
这个坑和 tenferro 卷的矩阵乘结果不对,也没有报错是一家:那里是 from_vec_col_major 把按行写的缓冲区读成按列,这里是 reshape 把按行理解的形状按列填。两处都只查元素个数。
修法
按行的 reshape 可以用列主序的 reshape 拼出来:把轴全部倒过来,reshape 成倒过来的目标形状,再把轴倒回去。
fn reshape_c(t: &TracedTensor, shape: &[usize]) -> R<TracedTensor> {
let rev_in: Vec<usize> = (0..t.rank).rev().collect();
let rev_shape: Vec<usize> = shape.iter().rev().copied().collect();
let rev_out: Vec<usize> = (0..shape.len()).rev().collect();
Ok(t.transpose(&rev_in)?.reshape(rev_shape)?.transpose(&rev_out)?)
}
数据本来就在手里时更简单:直接按行建矩阵。
let a2 = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_row_major(vec![2, 3], data.clone())?)?;
== fix 1: reshape_c (reverse axes, reshape, reverse back)
rows [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] a @ v = [210.0, 543.0]
reshape_c(reshape_c(arange(8), [2,2,2]), [4, 2]) rows [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]
...
== fix 2: build the matrix from row-major data
rows [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] a @ v = [210.0, 543.0]
两种都和 JAX 逐项相同(整数值,没有舍入)。reshape_c 多两个 transpose,在 traced 图里是两个运算;只在需要和 JAX 中间结果逐项对照时用它。整段程序自己一致时(例如 TRG 的并轴和拆轴都在 tenferro 里做),直接用列主序的 reshape,不必改,理由见布局与 reshape 的 TRG 一节。
怎么确认没中招
对 reshape 出来的矩阵抽查一个非对角元素,
a[0, 1]按 JAX 的习惯应该是 1:reshape: a[0, 1] = 2.0 reshape_c: a[0, 1] = 1.0 from_vec_row_major: a[0, 1] = 1.0列主序的
[i, j]在缓冲区里的位置是i + m·j。测试数据别用对称的或全相同的元素;
arange这种每个位置都不同的数据才查得出顺序。在 JAX 一侧临时改成
reshape(.., order="F")再跑:结果和 tenferro 一致,说明差别只在 reshape 顺序。
小结
- tenferro 的
reshape按列主序填,等于 JAX 的order="F";JAX 的默认是order="C"。照搬arange(n).reshape(..)会静默得到另一个矩阵。 - 按行的 reshape:倒轴、reshape 成倒过来的形状、再倒回来;或者用
from_vec_row_major直接建。 - 抽查一个非对角元素,比比较范数、迹这类不变量可靠。