unsupported dtype F32 for op index-select

主线位置：推理主线「前向」的第一步。token id 进 Embedding。

版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  lit["new 1 2 3"] --> i32["I32 下标"]
  i32 --> emb["Embedding forward"]
  emb -->|CPU| e1["报 F32 不支持"]
  emb -->|Metal| e2["报 I32 F32"]
  u8["U8 的 1 0 1"] --> sel["index_select"]
  sel --> three["取出 3 个元素"]
  lit -->|"写 1u32"| ok["U32 下标"]

本页的对照：PyTorch 侧是实验代码里的 p03.py，candle 侧是实验 crate pita-t 的 p03_i32（features metal accelerate，release）。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

测试时手写一组 token id，交给词嵌入：

emb = torch.nn.Embedding(4, 2)   # weight = arange(8).reshape(4, 2)
ids = torch.tensor([[1, 2, 3]])
emb(ids); emb(ids.int())
row = torch.tensor([10.0, 20.0, 30.0])
keep = torch.tensor([1, 0, 1], dtype=torch.uint8)
row[keep.bool()]

ids torch.int64 emb(ids) [[[2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]
emb(int32 ids) [[[2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]
row[keep.bool()] [10.0, 30.0]

整数字面量是 int64，int32 下标也行。0/1 的 uint8 张量转成 bool 就是掩码，选出第 0、2 个。

搬到 candle 的现象

let emb = Embedding::new(Tensor::arange(0f32, 8., &dev)?.reshape((4, 2))?, 2);
let ids = Tensor::new(&[[1, 2, 3]], &dev)?;
emb.forward(&ids);
let keep = Tensor::new(&[1u8, 0, 1], &dev)?;
row.index_select(&keep, 0);

== cpu
ids dtype I32
emb.forward(&ids)                  Err: unsupported dtype F32 for op index-select
row.index_select(&keep, 0)         F32 [3] [20.0, 10.0, 20.0]
== metal
ids dtype I32
emb.forward(&ids)                  Err: Metal contiguous index_select I32 F32 not implemented
row.index_select(&keep, 0)         F32 [3] [20.0, 10.0, 20.0]

两个现象：

• CPU 的报错说 F32 不支持，但 F32 是词表的 dtype，出问题的是下标；Metal 的报错把两个 dtype 都列了出来，更好查。
• U8 的 [1, 0, 1] 被当成下标，取出第 1、0、1 个元素，得到三个数，不报错。

原因

Tensor::new 按 Rust 的类型推断定 dtype，没有后缀的整数字面量是 i32，于是是 I32（PyTorch 是 int64）。Embedding::forward 把下标展平后调 index_select，CPU 后端的 index_select 按下标的 dtype 分派，只认 U8、U32、I64，其余落到最后一支：

            _ => Err(Error::UnsupportedDTypeForOp(self.dtype(), "index-select").bt()),

self 是被索引的词表，不是 ids，所以错误里写的是 F32。Metal 的 index_select 按（下标 dtype，数据 dtype）选内核，没有 I32 的组合，报 Metal contiguous index_select I32 F32 not implemented。Metal 上也不能先把 I32 转成 U32：to_dtype 报 Metal contiguous to_dtype I32 U32 not implemented（实测）。

U8 能当下标，是因为它在支持的三种下标 dtype 里；candle 没有布尔 dtype，比较运算返回的 U8 掩码也就和 U8 下标是同一种张量，index_select 无从区分。下标 dtype 的完整表格和 PyTorch 的对照见整数运算的「下标的 dtype」和归约与索引。

改法

下标写成 u32（tokenizers 的 get_ids() 本来就是 &[u32]）。掩码选择要先算出位置：

emb.forward(&Tensor::new(&[[1u32, 2, 3]], &dev)?);
ids.to_dtype(DType::U32).and_then(|i| emb.forward(&i)); // CPU only
let pos: Vec<u32> = keep.to_vec1::<u8>()?.iter().enumerate()
  .filter(|(_, k)| **k != 0).map(|(i, _)| i as u32).collect();
row.index_select(&Tensor::new(pos.as_slice(), &dev)?, 0);

fix: new(&[[1u32, 2, 3]])          F32 [1, 3, 2] [2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
fix: ids.to_dtype(U32)             F32 [1, 3, 2] [2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
fix: positions of keep             F32 [2] [10.0, 30.0]

第一行和第三行在 Metal 上相同。结果长度取决于数据的掩码选择，candle 没有对应的算子，只能在主机上算位置；如果下游接受「形状不变、不要的位置填一个值」，就用 where_cond（见shape mismatch in add / matmul / where_cond）。

怎么检查

• 读到 unsupported dtype <数据的 dtype> for op index-select（gather 同样），先看下标的 dtype，不是数据的。
• 测试代码里手写的下标一律带后缀：&[1u32, 2, 3]。从 numpy / JSON 读来的 id 用 Vec<u32> 接。
• 搬 x[mask] 时，确认 mask 是不是 0/1：是的话它在 candle 里不能直接当下标。

小结

• Tensor::new(&[1, 2, 3]) 是 I32，candle 的下标只收 U8、U32、I64；CPU 报 unsupported dtype F32 for op index-select（写的是词表的 dtype），Metal 报 index_select I32 F32 not implemented。
• U8 张量当下标就是下标，[1, 0, 1] 取出三个元素；PyTorch 的 x[mask.bool()] 要在主机上先算出位置。
• 改法：下标写 u32；CPU 上可以 to_dtype(DType::U32)，Metal 上 I32 转不了。