no method named topk found
主线位置：张量主线的「计算」一步，照着 transformers 的模型代码逐行翻译，碰到 PyTorch 常用、candle 没有的算子。
版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate）
flowchart TB py["PyTorch 算子"] --> tk["topk"] py --> mf["masked_fill"] py --> nz["nonzero"] py --> es["einsum"] tk --> a1["argsort gather"] mf --> a2["mask where_cond"] nz --> a3["读回主机自己数"] es --> a4["broadcast 或 matmul"]
实验代码：PyTorch 侧 ops.py，candle 侧 crate pitb-fail（会编译失败的那份）、pitb-core（只依赖 candle-core 的同一份），以及 crate pitb-exp 的 bin p25_ops（release）。平台 M2 Max、macOS 27、rustc 1.99.0。transformers 里这几个算子用得有多普遍，什么时候选 candle 数过。
PyTorch 的写法
x = torch.tensor([[0.1, 0.9, 0.3, 0.7], [0.5, 0.2, 0.8, 0.4]])
v, i = torch.topk(x, 2, dim=-1)
mask = torch.tensor([[False, True, False, True], [True, False, False, False]])
print("masked_fill =", x.masked_fill(mask, float("-inf")).tolist())
print("nonzero =", mask.nonzero().tolist())
print("einsum =", torch.einsum("i,j->ij", torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0, 5.0])).tolist())
q = torch.arange(24, dtype=torch.float32).reshape(1, 2, 3, 4) / 24
print("einsum bhqd,bhkd->bhqk [0,1,2] =", torch.einsum("bhqd,bhkd->bhqk", q, q)[0, 1, 2].tolist())
topk values = [[0.8999999761581421, 0.699999988079071], [0.800000011920929, 0.5]] indices = [[1, 3], [2, 0]]
masked_fill = [[0.10000000149011612, -inf, 0.30000001192092896, -inf], [-inf, 0.20000000298023224, 0.800000011920929, 0.4000000059604645]]
nonzero = [[0, 1], [0, 3], [1, 0]]
einsum = [[3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]
einsum bhqd,bhkd->bhqk [0,1,2] = [2.024305582046509, 2.621527671813965, 3.21875]
搬到 candle 的现象
照着写：
let _top = x.topk(2)?;
let _filled = x.masked_fill(&mask, f32::NEG_INFINITY)?;
let _nz = mask.nonzero()?;
let _outer = candle_core::einsum("i,j->ij", &[&x, &x])?;
在一个同时依赖 candle-transformers 的 crate 里（pitb-fail，cargo build --color never；只截取 einsum 和 topk 两条，指向 cargo 缓存里 deepseek2.rs 的那一行换成了 ...）：
error[E0425]: cannot find function `einsum` in crate `candle_core`
--> src/bin/missing.rs:10:29
|
10 | let _outer = candle_core::einsum("i,j->ij", &[&x, &x])?;
| ^^^^^^ not found in `candle_core`
error[E0599]: no method named `topk` found for struct `candle_core::Tensor` in the current scope
--> src/bin/missing.rs:7:16
|
7 | let _top = x.topk(2)?;
| ^^^^ method not found in `candle_core::Tensor`
|
...
109 | fn topk(&self, topk: usize) -> Result<TopKOutput>;
| ---- the method is available for `candle_core::Tensor` here
|
= help: items from traits can only be used if the trait is in scope
help: trait `TopKLastDimOp` which provides `topk` is implemented but not in scope; perhaps you want to import it
|
1 + use candle_transformers::models::deepseek2::TopKLastDimOp;
|
masked_fill 是一条不带 help 的 E0599，nonzero 和 topk 一样提示 candle_transformers::models::deepseek2::NonZeroOp。同一份代码放进只依赖 candle-core 的 crate（pitb-core），四条错误还在，两条 help 都没了：
error[E0599]: no method named `topk` found for struct `candle_core::Tensor` in the current scope
| ^^^^ method not found in `candle_core::Tensor`
原因
candle-core 0.11.0 没有 topk、masked_fill、nonzero、einsum，origin/main 上也没有。rustc 的提示来自 DeepSeek V2 的模型文件：它为自己用到的两个算子定义了 TopKLastDimOp 和 NonZeroOp，并为 Tensor 实现，rustc 在依赖里找到了就建议导入。这不是 candle 推荐的 API，只是某个模型的内部工具碰巧是 pub。masked_fill 则是每个模型文件各写一份私有函数：candle-transformers 0.11.0 里有 23 个文件定义了自己的 fn masked_fill（llama.rs 里的见对齐 logits）。
deepseek2 的 topk 就是三步：
fn topk(&self, topk: usize) -> Result<TopKOutput> {
// Sorted descending
let sorted_indices = self.arg_sort_last_dim(false)?;
let topk_indices = sorted_indices.narrow(D::Minus1, 0, topk)?.contiguous()?;
Ok(TopKOutput {
values: self.gather(&topk_indices, D::Minus1)?,
indices: topk_indices,
})
}
nonzero 的输出形状取决于数据，没法用已有的算子拼，NonZeroOp 是一个只有 CPU 实现的自定义算子，别的设备上先拷回 CPU。
改法
/// torch: x.masked_fill(mask, value). candle: mask.where_cond(on_true, on_false).
fn masked_fill(x: &Tensor, mask: &Tensor, value: f32) -> candle_core::Result<Tensor> {
let v = Tensor::new(value, x.device())?.to_dtype(x.dtype())?.broadcast_as(x.shape())?;
mask.where_cond(&v, x)
}
let idx = x.arg_sort_last_dim(false)?.narrow(D::Minus1, 0, 2)?.contiguous()?;
println!("topk values = {:?} indices = {:?}", x.gather(&idx, D::Minus1)?.to_vec2::<f32>()?, idx.to_vec2::<u32>()?);
println!("masked_fill = {:?}", masked_fill(&x, &mask, f32::NEG_INFINITY)?.to_vec2::<f32>()?);
println!("nonzero = {:?}", nonzero_2d(&mask)?); // to_vec2 读回主机，逐个数
// einsum("i,j->ij")：外积
println!("einsum i,j->ij = {:?}", a.unsqueeze(1)?.broadcast_mul(&b.unsqueeze(0)?)?.to_vec2::<f32>()?);
// einsum("bhqd,bhkd->bhqk")：注意力分数
println!("einsum bhqd,bhkd->bhqk [0,1,2] = {:?}", q.matmul(&q.t()?)?.get(0)?.get(1)?.get(2)?.to_vec1::<f32>()?);
topk values = [[0.9, 0.7], [0.8, 0.5]] indices = [[1, 3], [2, 0]]
deepseek2 topk values = [[0.9, 0.7], [0.8, 0.5]] indices = [[1, 3], [2, 0]]
masked_fill = [[0.1, -inf, 0.3, -inf], [-inf, 0.2, 0.8, 0.4]]
where_cond swapped = [[-inf, 0.9, -inf, 0.7], [0.5, -inf, -inf, -inf]]
nonzero = [[0, 1], [0, 3], [1, 0]]
deepseek2 nonzero = [[0, 1], [0, 3], [1, 0]] U32
einsum i,j->ij = [[3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]
einsum bhqd,bhkd->bhqk [0,1,2] = [2.0243056, 2.621528, 3.2187502]
- 结果和 PyTorch 一致（打印精度不同；
einsum那一行差在最后一位，是 f32 求和顺序）。导入deepseek2的两个 trait 也能得到相同的结果，但把一个模型文件的内部工具当成公共 API 依赖，换版本时可能消失，自己写这几行更稳。 where_cond的参数顺序和masked_fill相反。x.masked_fill(mask, v)是「mask 为真的地方填 v」；candle 是mask.where_cond(on_true, on_false)，填充值放第一个。where_cond swapped一行是把x和-inf写反的结果：保留的和屏蔽的正好颠倒，不报错。- 索引的 dtype：
topk的下标、nonzero都是 U32，PyTorch 是 int64（归约与索引）。 arg_sort_last_dim只排最后一维，topk(dim=1)之类要先transpose。einsum没有通用的替代，按下标式拆成broadcast_mul+sum、matmul、transpose的组合，每一种都要对一次数。
怎么检查
- 翻译模型时先列出 PyTorch 代码里用到的算子（
grep -o "\.\w*(" modeling_x.py | sort -u），对着 candle 的Tensor文档逐个找对应，没有的提前写好替代函数和测试，不要等编译器一个个报。 where_cond的替代函数配一个两元素的单元测试（一个 mask 为真、一个为假），参数写反立刻能看出来。- rustc 提示
use candle_transformers::models::<某个模型>::...时，这是在建议依赖一个模型的私有工具，先看一眼它的实现，通常几行就能自己写。
小结
topk、masked_fill、nonzero、einsum在 candle-core 0.11.0 里都没有（main 上也没有），编译报 E0599 / E0425。- rustc 建议导入
candle_transformers::models::deepseek2::TopKLastDimOp/NonZeroOp，只因为 crate 依赖了 candle-transformers；只依赖 candle-core 时没有这条提示。 - 替代：
topk= 降序arg_sort_last_dim+narrow+gather；masked_fill=mask.where_cond(填充值, x)；nonzero读回主机自己数；einsum按式子拆成广播乘和matmul。和 PyTorch 结果一致。 where_cond的填充值在第一个参数，写反不报错、结果颠倒。
💬 关联讨论
发起讨论加载中...