布局与 reshape

主线位置：张量主线的第一步和第三步，建数组时一段平铺数据怎么解释、读回主机时缓冲区按什么顺序排；顺带 TRG 主线第二步里那次「transpose 再 reshape」。

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

flowchart TB
  flat["平铺数据 0 到 5"] -->|"jnp reshape"| jc["C 序 最后一维最快"]
  flat -->|"from_vec_col_major"| tc["列主序 第一维最快"]
  flat -->|"from_vec_row_major"| tr["重排一次"]
  tr --> tc2["列主序 值同 C 序"]
  tc -->|"reshape"| tf["等于 order F"]
  tc -->|"as_slice"| buf["列主序缓冲区"]
  buf -->|"reshape order F"| npy["NumPy 数组"]

实验代码：一个 JAX 脚本（写出随机张量、跑一步 TRG）、一个 tenferro 程序 layout（release，CpuEngine::serial()），再用一个 NumPy 脚本把 tenferro 写回的缓冲区读回来对比。平台 M2 Max、macOS 27、rustc 1.99.0。tenferro 只有列主序这一种拥有型布局、视图怎么带步长，见 tenferro 卷的列主序与视图；这里只讲从 JAX 搬代码时会碰到的几处。

同一段数据，两种读法

把 0, 1, ..., 5 当成 2×3 矩阵，乘 v = [1, 10, 100]：

data = jnp.arange(6.0)
v = jnp.array([1.0, 10.0, 100.0])
a_c = data.reshape(2, 3)
a_f = jnp.reshape(data, (2, 3), order="F")
print(a_c @ v, a_f @ v)

let data: Vec<f64> = (0..6).map(|i| i as f64).collect();
let a_cm = Tensor::from_vec_col_major(vec![2, 3], data.clone())?;
let a_rm = Tensor::from_vec_row_major(vec![2, 3], data.clone())?;
// matvec: TracedTensor::dot_general with v, contracting a's axis 1

# JAX
reshape(2, 3)           rows [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]  A @ v = [210.0, 543.0]
reshape(.., order='F')  rows [[0.0, 2.0, 4.0], [1.0, 3.0, 5.0]]  A @ v = [420.0, 531.0]
# tenferro
from_vec_col_major rows [[0.0, 2.0, 4.0], [1.0, 3.0, 5.0]]  A @ v = [420.0, 531.0]
from_vec_row_major as_slice [0.0, 3.0, 1.0, 4.0, 2.0, 5.0]
from_vec_row_major rows [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]  A @ v = [210.0, 543.0]

JAX 和 NumPy 默认 C 序：最后一维变得最快，[i, j] 在平铺数据里的位置是 3i + j。
tenferro 的 from_vec_col_major 按列主序读：第一维变得最快，位置是 i + 2j。同一个 Vec 读出来是另一个矩阵，形状一样，所以不报错。它和 JAX 的 order="F" 是同一种读法，结果 [420, 531] 也一样。
数据本来就是 C 序（从 NumPy、从 C 库、从按行写的嵌套 Vec 来）时，用 from_vec_row_major。它重排一次，得到的仍是列主序张量：as_slice 是 [0, 3, 1, 4, 2, 5]，值和 JAX 的 a_c 相同。这次复制的开销和别的桥接办法见列主序与视图；交给 from_vec_col_major 会怎样，见常见坑矩阵乘结果不对，也没有报错。

reshape 也按列主序

reshape 不移动数据，只换一种读法。两边都是这样，所以 tenferro 的 reshape 继承了列主序：

x = jnp.arange(8.0).reshape(2, 2, 2)
print(x.reshape(4, 2))
xf = jnp.reshape(jnp.arange(8.0), (2, 2, 2), order="F")
print(jnp.reshape(xf, (4, 2), order="F"))

let x = TracedTensor::from_vec_col_major(vec![2, 2, 2], (0..8).map(|i| i as f64).collect())?;
let y = eval(&x.reshape([4, 2])?)?;

# JAX
C order: reshape(4,2) rows [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]
F order: reshape(4,2,'F') rows [[0.0, 4.0], [1.0, 5.0], [2.0, 6.0], [3.0, 7.0]]
# tenferro
reshape([4, 2]) rows [[0.0, 4.0], [1.0, 5.0], [2.0, 6.0], [3.0, 7.0]]

TracedTensor::reshape 的文档写的是「Reshape without changing element order」，元素顺序就是列主序那个顺序。所以 tenferro 的 x.reshape(s) 对应 jnp.reshape(x, s, order="F")，不对应 x.reshape(s)。tenferro 没有 order 参数。

TRG 的 reshape 为什么不用改

TRG 每一步把四脚张量 T[u, l, d, r] 换轴、并成 d² × d² 矩阵做 SVD，再把奇异向量拆回三脚张量。这几行在两边写得一模一样：

ma = t.transpose(2, 1, 0, 3).reshape(d * d, d * d)
s1 = (ua[:, :k] * jnp.sqrt(sa[:k])).reshape(d, d, k)

let ma = t.transpose(&[2, 1, 0, 3])?.reshape([m, m])?;
let left = (&cols(&u)? * &skb)?.reshape([d, d, k])?;

两边并出来的矩阵不一样。JAX 把 (a, b) 并成行号 d·a + b，tenferro 并成 a + d·b。把同一个随机张量（[3, 3, 3, 3]，NumPy 种子 0）交给两边各跑一步（k = 5），JAX 侧用 C 序和 order="F" 各跑一遍：

# JAX
t[1,0,2,1] = np.float64(0.2146591225063409)
order=C: ma[0,:3] = [0.1257302210933933, -0.1321048632913019, 0.6404226504432821]
order=C: s[:5] = [4.927039373327931, 4.347056795845394, 3.759268211999318, 2.9569278616680923, 2.1963353167288187]
order=C: trace(T') = 5.217301654602105  |T'| = 33.63471589099613
order=F: ma[0,:3] = [0.1257302210933933, -0.9217253762584194, -1.2883614637495544]
order=F: s[:5] = [4.927039373327933, 4.347056795845394, 3.7592682119993173, 2.956927861668093, 2.196335316728822]
order=F: trace(T') = 5.217301654602104  |T'| = 33.63471589099611
...
max| |T'_C| - |T'_F| | elementwise = 4.1744385725905886e-14
# tenferro
T_f.bin (ravel order='F'): t[1,0,2,1] = 0.2146591225063409
T_f.bin (ravel order='F'): ma[0,:3] = [0.1257302210933933, -0.9217253762584194, -1.2883614637495544]
T_f.bin (ravel order='F'): s[:5] = [4.927039373327933, 4.347056795845393, 3.7592682119993133, 2.9569278616680923, 2.19633531672882]
T_f.bin (ravel order='F'): trace(T') = 5.217301654602115  |T'| = 33.63471589099614
...
# tenferro 写回的 ma、T' 用 NumPy 按 order='F' 读回后对比 JAX 的 order=F
ma: tenferro vs jnp.reshape(order='F') bitwise equal: True
...
T': max| |T'_tf| - |T'_jax| | = 3.505529200253932e-14
T': elements with opposite sign: 272 of 625

tenferro 的矩阵 ma 和 JAX order="F" 的逐位相同，和 JAX 默认的 C 序不同：第一行只有第一个数一样。
奇异值、trace(T')（T'[x, y, x, y] 的和，lnZ 的最后一项就是它的对数）、‖T'‖ 都对得上。tenferro 对 JAX 的 C 序：奇异值最大差 4.9e-15（量级 4.9），trace(T') 差 1.0e-14（量级 5.2）。不是 0，是因为两边的 SVD 实现不同（jaxlib 走 LAPACK，tenferro 走 faer），收缩的求和顺序也不同。
T' 本身有 272 个元素符号相反，绝对值差 3.5e-14。这是 SVD 的符号自由度：每一对奇异向量 (u_i, v_i) 可以同时取反，两个库的取法不同，T' 的新指标随之差一个 ±1。lnZ 只用到和这个符号无关的量，所以不受影响。

并法不同为什么不影响结果？换一种并法只是给矩阵的行和列重新编号：M_F = P · M_C · Q，P、Q 是置换矩阵。SVD 随之变成 (P·U) · S · (Qᵀ·V)ᵀ，奇异值不变，奇异向量的行被同样置换；拆回 [d, d, k] 时用的是同一种约定，置换被还原。只要一段程序里所有的并和拆都用同一种约定，结果就是同一个张量。完整的 20 步 TRG 两边的对照见动手：Ising TRG。

约定只在越过边界时才要紧：数据从文件或 NumPy 进来、结果交回 Python，或者平铺下标本身有含义（例如格点 (x, y) 编号成 x + L·y 还是 L·x + y）。

❗IMPORTANT

动脑筋：把 NumPy 自己的 C 序字节（t.ravel(order="C")）直接交给 from_vec_col_major(vec![3, 3, 3, 3], ..)，再跑同样一步 TRG，trace(T') 是多少？

形状是 [3, 3, 3, 3]，元素个数对得上，不会报错。想一想按列主序读 C 序字节得到的是哪个张量，再想一想 T'[x, y, x, y] 这个求和对它敏感不敏感。

读回主机

tenferro 的 as_slice、host_data 给出的是列主序缓冲区。交回 NumPy 时要按列主序解释：

ma_tf = np.fromfile("tf_ma.bin").reshape(9, 9, order="F")

write_f64("tf_ma.bin", ma.as_slice::<f64>()?)?;  // raw little-endian f64

ma: tenferro vs jnp.reshape(order='F') bitwise equal: True
ma: read back with order='C' equal: False

反方向同理：NumPy 数组交给 tenferro 之前写成 x.ravel(order="F")，在 Rust 侧用 from_vec_col_major；拿到的是 C 序缓冲区时用 from_vec_row_major。np.asarray(jax_array) 总是 C 序（实测 C_CONTIGUOUS True），ravel() 也按 C 序走。

批维在右边

JAX 的批量线代把批维放在最左边：jnp.linalg.svd 对 (4, 2, 3) 做 4 个 2×3 矩阵。tenferro 的约定是 [M, N, B...]，批维在右边，每个批次在列主序下占一段连续内存（为什么这样定，见列主序与视图）。同样 4 个矩阵：

xs = rng.standard_normal((4, 2, 3))
np.moveaxis(xs, 0, -1).ravel(order="F").tofile("X_b.bin")   # tenferro: [2, 3, 4]
s = jnp.linalg.svd(jnp.asarray(xs), compute_uv=False)

let xb = Tensor::from_vec_col_major(vec![2, 3, 4], read_f64("X_b.bin")?)?;
let (_, sb, _) = TracedTensor::from_tensor_concrete_shape(xb)?.svd()?;

# JAX
== batched SVD: JAX x.shape (4, 2, 3) -> s.shape (4, 2)
s[:, 0] = [1.6788972654264325, 0.8761945108602387, 0.9888740800975825, 1.6642548983283698]
# tenferro
x.shape [2, 3, 4] -> s.shape [2, 4]
s[0, :] = [1.6788972654264323, 0.8761945108602388, 0.9888740800975825, 1.66425489832837]
x.shape [4, 2, 3] (JAX order) -> s.shape: [2, 3]
batched s: max|s_tf.T - s_jax| = 4.440892098500626e-16  max s = 1.6788972654264325

批维挪到右边以后，两边的奇异值最大差 4.4e-16（量级 1.7），s 的形状也是批维在右：[2, 4] 对 JAX 的 (4, 2)。
把 JAX 形状 [4, 2, 3] 原样交给 tenferro，它把前两维当成矩阵、最后一维当批：3 个 4×2 矩阵，s 是 [2, 3]，没有任何错误。
dot_general 的输出轴序也是批维在最后（[lhs free..., rhs free..., batch...]），和 lax.dot_general 的批维在前相反，见 einsum。

📝NOTE

没有蠢问题

问： JAX 也有 order="F"。把 JAX 侧所有的 reshape 都改成 order="F"，两边的输出就逐位一致了吧？

答：只有布局那一步能做到。上面的实验里 ma 逐位相同，但奇异值仍差 4.0e-15、T' 仍有一半元素符号相反：SVD 由两个不同的库算，求和顺序和符号约定都不同。布局对上以后剩下的差别是数值实现的差别，判断标准见怎么读这一卷。改 order 的价值在调试：它让两边的中间张量可以逐项比较，而不是只能比较不变量。

动脑筋答案

T_c.bin (ravel order='C'): t[1,0,2,1] = 1.4580206835369587
T_c.bin (ravel order='C'): ma[0,:3] = [0.1257302210933933, -0.1321048632913019, 0.6404226504432821]
T_c.bin (ravel order='C'): s[:5] = [4.927039373327933, 4.347056795845393, 3.7592682119993133, 2.9569278616680923, 2.196335316728821]
T_c.bin (ravel order='C'): trace(T') = 5.217301654602139  |T'| = 33.634715890996134

张量读错了：t[1,0,2,1] 应该是 0.2146591225063409（JAX 侧和 T_f.bin 都是它），这里是 1.4580206835369587。
但 trace(T') 是 5.217301654602139，和正确的 5.217301654602115 只差 2.4e-14，奇异值也相同。
原因：C 序把最后一维放在最快的位置，列主序把第一维放在最快的位置。C 序的字节按列主序读，得到的是把所有轴倒过来的张量 t_rev[i, j, k, l] = t[l, k, j, i]。再用列主序并轴，又正好等于 C 序的并法：ma 第一行和 JAX order=C 的那一行逐位相同。整步计算变成了 JAX 默认 C 序那次计算，只是 T' 的轴也倒了过来；而 T'[x, y, x, y] 的求和把轴倒过来还是它自己。
所以 lnZ、奇异值这类不变量查不出布局错误。读数据之后至少抽查一个元素。

小结

JAX / NumPy 默认 C 序（最后一维最快），tenferro 只有列主序（第一维最快）；同一个 Vec 交给 from_vec_col_major 等于 JAX 的 order="F"，C 序数据用 from_vec_row_major。
tenferro 的 reshape 按列主序，对应 jnp.reshape(.., order="F")，没有 order 参数。
一段程序里并轴和拆轴用同一种约定时，结果和约定无关：TRG 的 transpose 加 reshape 原样搬过来，奇异值、trace(T') 和 JAX 差 1e-14 量级；T' 本身差 SVD 的符号。
边界上要写明顺序：读回用 reshape(shape, order="F")，送进去用 ravel(order="F")；C 序字节当列主序读等于把轴全部倒过来，不变量可能完全看不出来。
批量线代的批维在右边（[M, N, B]），JAX 形状原样传进来不会报错，只会算另一批矩阵。