unexpected rank, expected: 0, got: 2

主线位置：算子主线的读回一步和推理主线的取样一步。to_scalar、to_vec1、to_vec2 把张量读成 Rust 值之前先核对秩；LogitsProcessor::sample 内部用的正是它们，见取样。

flowchart TB
  t["形状 1 x 1 的 loss"] --> rd["to_scalar"]
  rd --> chk{"rank == 0"}
  chk -->|否| err["unexpected rank"]
  fix["squeeze 或 sum_all"] --> r0["0 维张量"]
  r0 --> ok["f32"]

实验 crate exp-candle/pa-pit，bin p05_rank，CPU，debug；M2 Max、macOS 27、rustc 1.99.0。

现象

两个常见来源。一是归约时留了维度：

let a = Tensor::ones((2, 3), DType::F32, &dev)?;
let loss = a.sum_keepdim(0)?.sum_keepdim(1)?;
println!("loss.to_scalar: {:?}", loss.to_scalar::<f32>().map_err(|e| e.to_string()));

loss dims [1, 1]
loss.to_scalar: Err("unexpected rank, expected: 0, got: 2 ([1, 1])")
a.to_vec1: Err("unexpected rank, expected: 1, got: 2 ([2, 3])")
a.mean(1).to_vec2: Err("unexpected rank, expected: 2, got: 1 ([2])")

只有一个元素也不行，[1, 1] 的秩是 2。最后一行反过来：mean(1) 不带 keepdim，维度已经去掉了。

二是模型输出带 batch 维。quantized_llama::ModelWeights::forward 返回 (batch, vocab)，忘了 squeeze(0) 就交给取样：

let logits = Tensor::new(&[[2.0f32, 1.5, 1.0, 0.5, 0.0, -1.0]], &dev)?;
let mut lp = LogitsProcessor::from_sampling(1, s.clone());
lp.sample(&logits)

ArgMax on [1, 6]: Err("unexpected rank, expected: 0, got: 1 ([1])")
All { temperature: 0.8 } on [1, 6]: Err("unexpected rank, expected: 1, got: 2 ([1, 6])")
GumbelSoftmax { temperature: 0.8 } on [1, 6]: Err("unexpected rank, expected: 0, got: 1 ([1])")
apply_repeat_penalty on [1, 6]: Err("unexpected rank, expected: 1, got: 2 ([1, 6])")

ArgMax 和 GumbelSoftmax 报的形状是 [1]，不是传进去的 [1, 6]，从报错很难联想到 logits 多了一维。

出处

Tensor::to_scalar 开头：

        if self.rank() != 0 {
            Err(Error::UnexpectedNumberOfDims {
                expected: 0,
                got: self.rank(),
                shape: self.shape().clone(),
            }
            .bt())?
        }

to_vec1、to_vec2、to_vec3 同理，分别要求秩 1、2、3。错误里的形状是调用 to_* 的那个张量的形状。

LogitsProcessor::sample 的几种取法读值的地方不同：

ArgMax：logits.argmax(D::Minus1)?.to_scalar::<u32>()。(1, 6) 的 argmax 沿最后一维得到 (1)，to_scalar 看到的是 [1]。
GumbelSoftmax：同样是 argmax 之后 to_scalar。
All、TopK、TopP：先 softmax_last_dim，再 to_vec1()，看到的是 [1, 6]。

apply_repeat_penalty 第一步也是 to_vec1。

修法

读值之前把维度去掉。按场景选：

let loss = loss.squeeze(0)?.squeeze(0)?;   // drop known size-1 dims
let total = a.sum_all()?;                    // reduce everything to rank 0
let flat = a.flatten_all()?.to_vec1::<f32>()?;
let logits = model.forward(&input, pos)?.squeeze(0)?; // (1, vocab) -> (vocab)

loss.squeeze(0)?.squeeze(0)?.to_scalar: 6.0
a.sum_all()?.to_scalar: 6.0
a.flatten_all()?.to_vec1: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
a.mean(1)?.to_vec1: [1.0, 1.0]
ArgMax on squeeze(0) [6]: 0
ArgMax on i(0) [6]: 0
argmax(D::Minus1) on (1,6): [0]

损失这类标量：归约用 sum_all / mean_all，结果直接是 0 维。
模型输出：candle-examples 里的 quantized 例子就是 model.forward(&input, 0)? 之后紧跟 logits.squeeze(0)?。squeeze(0) 在第 0 维长度不为 1 时原样返回（不报错），批大小大于 1 时要改用 i(b) 逐条取。
批量取样不要交给 LogitsProcessor，它只处理一行；argmax(D::Minus1)?.to_vec1::<u32>() 一次拿到每一行的结果（上面最后一行）。

怎么确认没中招

报错里的 expected 是你调用的 to_* 的秩（to_scalar 是 0，to_vec1 是 1），got 和括号里的形状属于调用它的那个张量；在库函数里报出来时，先想它在内部对你的输入做了什么（argmax、softmax）。
在模型输出和取样之间加一行 assert_eq!(logits.rank(), 1)。

小结

to_scalar / to_vec1 / to_vec2 严格按秩读值，元素个数为 1 也不行。
keepdim 归约和模型输出的 batch 维是两大来源；LogitsProcessor::sample 只吃一维 logits。
修法：squeeze、sum_all、flatten_all、i(..)；批量取样直接 argmax。

unexpected rank, expected: 0, got: 2

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录