lu_factor_solve: rank mismatch: expected 2, actual 1

主线位置：张量主线的「计算」：线性方程组 A x = b。

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

flowchart TB
  b1["b 形状 5"] --> jx["jnp.linalg.solve"]
  jx --> x1["x 形状 5"]
  b1 --> tf["a.solve b"]
  tf --> err["rank mismatch"]
  b1 --> rs["reshape 成 5 1"]
  rs --> tf2["a.solve"]
  tf2 --> back["reshape 回 5"]

实验代码：tenferro 一侧是 bin p23_solve_1d（release，CpuEngine::serial()），JAX 一侧是 p23_solve_1d.py。M2 Max、macOS 27、rustc 1.99.0。

现象

A = 3I + M，M[i, j] = sin(1 + i + 2j)，b[i] = cos(i)，损失 L = w · x。JAX 里右端项是向量：

x = jnp.linalg.solve(A, b)
jax.grad(lambda A: jnp.sum(w * jnp.linalg.solve(A, b)))(A)

jax 0.12.0.1 x shape (5,)
x = [ 0.30143657  0.18361373 -0.10302273 -0.29494057 -0.21569141]
dL/dA[0, :] = [-0.13782667 -0.08395421  0.04710537  0.13485649  0.09862118]

tenferro 照搬：

let b = TracedTensor::from_vec_col_major(vec![5], bv.clone())?;
a.solve(&b)                                    // traced
ert.with_eager_session(&mut eng, |s| s.solve(&ea, &eb))   // eager

traced a.solve(&b[5]): Err(tenferro-linalg.lu_factor_solve: rank mismatch: expected 2, actual 1)
eager solve(a, b[5]): Err(solve: rank mismatch: expected 2, actual 1)

traced 在建图那一行就返回错误，报的是内部算子名 lu_factor_solve，不是你调用的 solve；eager 报的是 solve。

出处

traced 的 solve 建一个融合的 LuFactorSolve 算子（一次分解加求解，保存 LU 因子给反向用）。它的形状推导先检查右端项是不是矩阵（extension.rs）：

            LinalgOp::LuFactorSolve => {
                require_matrix_meta("tenferro-linalg.lu_factor_solve", input_shapes[1])?;
                // ...
                vec![
                    (promote_dtypes(&input_dtypes), input_shapes[1].to_vec()),

require_matrix_meta 要求秩至少为 2；输出的形状直接取 b 的形状，所以 tenferro 没有 NumPy「一维的 b 当一列、结果再压回一维」的特例。其余分解的梯度对照见分解的梯度。

修法

右端项 reshape 成一列，解再 reshape 回来：

let x = a.solve(&b.reshape([5, 1])?)?.reshape([5])?;
let l = (&x * &w)?.reduce_sum(None)?;
let g = ad.grad(&l, &a)?;

x = [0.3014365729163191, 0.18361373020203742, -0.10302272928189653, -0.294940566577713, -0.21569140715009136]
dL/dA[0, 0..5] = [-0.137826673718457, -0.0839542111229163, 0.04710536600435243, 0.13485648686474488, 0.0986211756242439]
max|jax - tenferro| = 8.327e-17, max|grad| = 0.231

两个 reshape 都在图里，可以求导；梯度和 JAX 差 8.3e-17（梯度最大 0.23，不到 1 个 ulp）。

怎么确认没中招

搜 solve(：右端项是一维的地方都会在建图时失败，不会静默。
要小心的是反方向：把 [n, 1] 的解直接当 [n] 用，和另一个 [n] 张量逐元素运算时形状不符。解出来马上 reshape 回一维。

小结

jnp.linalg.solve 接受一维 b；tenferro 的 solve 要求 b 至少二维，traced 报 lu_factor_solve: rank mismatch: expected 2, actual 1，eager 报 solve: rank mismatch。
写成 a.solve(&b.reshape([n, 1])?)?.reshape([n])?，梯度和 JAX 一致。

lu_factor_solve: rank mismatch: expected 2, actual 1

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录