Rust in Everything
基于 tenferro-runtime 0.8.0

同一段数据 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 直接建。
  • 抽查一个非对角元素,比比较范数、迹这类不变量可靠。

💬 关联讨论

发起讨论
加载中...