归约与索引
主线位置：张量主线的中间一步「计算」：按维归约、取下标，以及按条件、按下标张量取值。最后一步读回主机时，归约留下的秩决定了该用
to_scalar还是to_vec1。
版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate；同一组输入另在 Metal 上跑了一遍，不同之处单独标出）
flowchart TB t["Tensor 2x3"] -->|"sum 1"| r1["形状 2"] t -->|"sum_keepdim 1"| r2["形状 2x1"] t -->|"argmax 1"| ai["U32 形状 2"] t -->|"i 区间"| nv["narrow 视图"] t -->|"index_select ids"| cp["新张量"] ids["ids U8 U32 I64"] --> cp nv -.->|"越界"| e["Err"] cond["cond U8"] -->|"where_cond"| w["同形状结果"]
本页的对照：PyTorch 侧是实验代码里的 reduce.py 和 red_ref.py（torch.set_num_threads(1)），candle 侧是实验 crate tensora-t 的 reduce、red、intlit、neg_ids、gather_oob（features metal accelerate，release；reduce metal 在 Metal 上跑同一段代码）。平台 M2 Max、macOS 27、rustc 1.99.0。下面的 a 是 arange(6) 排成的 [2, 3] 浮点矩阵。
按维归约：去掉这一维，和 PyTorch 一样
a.sum(1); a.sum(1, keepdim=True); a.sum((0, 1)); a.sum(); a.mean(0)
a.sum(-1); a.sum(2); a.sum().item()
a.sum(1); a.sum_keepdim(1); a.sum((0, 1)); a.sum_all(); a.mean(0); a.mean_all();
a.sum(D::Minus1); a.sum(2); a.sum_all()?.to_scalar::<f32>();
a.sum(1) [2] torch.float32 [3.0, 12.0]
a.sum(1, keepdim=True) [2, 1] torch.float32 [3.0, 12.0]
a.sum((0, 1)) [] torch.float32 [15.0]
a.sum() [] torch.float32 [15.0]
a.mean(0) [3] torch.float32 [1.5, 2.5, 3.5]
a.sum(-1) [2] torch.float32 [3.0, 12.0]
a.sum(2) IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)
a.sum().item() 15.0
a.sum(1) [2] F32 [3.0, 12.0]
a.sum_keepdim(1) [2, 1] F32 [3.0, 12.0]
a.sum((0, 1)) [] F32 [15.0]
a.sum_all() [] F32 [15.0]
a.mean(0) [3] F32 [1.5, 2.5, 3.5]
a.mean_all() [] F32 [2.5]
a.sum(D::Minus1) [2] F32 [3.0, 12.0]
a.sum(2) Err: sum: dimension index 2 out of range for shape [2, 3]
a.sum_all()?.to_scalar::<f32>() Ok(15.0)
形状两边一致：sum 去掉被归约的维度，keepdim=True 换成另一个方法 sum_keepdim；两者都接受一个维度或一个维度元组，内部是同一个 sum_impl，最后按 keepdim 决定要不要 squeeze_dims。全部求和不是不带参数的 sum()，而是 sum_all / mean_all，结果是 0 维张量，to_scalar::<f32>() 对应 .item()。负的维度写 D::Minus1。
burn 在这里和两边都不同：它的 sum_dim 永远保留维度，全局 sum 返回 [1]，见 burn 卷的归约与广播。从 PyTorch 搬到 candle，归约这一层的形状不用改。
max 只给值，argmax 是 U32
a.max(1).values; a.max(1).indices; a.argmax(1)
ties = torch.tensor([[1.0, 3, 3], [2, 2, 0]])
ties.argmax(1); ties.argmin(1)
a.max(1); a.max_keepdim(1); a.argmax(1); a.argmax_keepdim(1);
let ties = Tensor::new(&[[1f32, 3., 3.], [2., 2., 0.]], &dev)?;
ties.argmax(1); ties.argmin(1);
a.max(1).values [2] torch.float32 [2.0, 5.0]
a.max(1).indices [2] torch.int64 [2, 2]
a.argmax(1) [2] torch.int64 [2, 2]
ties.argmax(1) [2] torch.int64 [1, 0]
ties.argmin(1) [2] torch.int64 [0, 2]
a.max(1) [2] F32 [2.0, 5.0]
a.max_keepdim(1) [2, 1] F32 [2.0, 5.0]
a.argmax(1) [2] U32 [2.0, 2.0]
a.argmax_keepdim(1) [2, 1] U32 [2.0, 2.0]
ties.argmax(1) [2] U32 [1.0, 0.0]
ties.argmin(1) [2] U32 [0.0, 2.0]
torch.max(x, dim) 返回值和下标两样；candle 的 max 只返回值，下标要另调 argmax，两次各遍历一遍。并列时两边都取第一个。
下标的 dtype 是 U32，不是 int64：CPU 后端的 ReduceIndex 把结果写成 CpuStorage::U32(..)，写死在代码里，没有参数。U32 正好能直接喂给 index_select、embedding（见下文），也是 LogitsProcessor 里 to_scalar::<u32>() 读的类型。拿它和 i64 的标签比较（算准确率）之前，先 to_dtype(DType::I64)，否则是 dtype mismatch。
let nan = Tensor::new(&[[1f32, f32::NAN, 2.], [f32::NAN, 0., 5.]], &Device::Cpu)?;
nan.max(1)?; nan.argmax(1)?;
PyTorch 里 nan.max(1).values 是 [nan, nan]，argmax(1) 是 [1, 0]：任何一行有 NaN，结果就是 NaN。candle CPU 上这两行各是多少？
var、整数的 mean
a.var(1)
torch.tensor([[1.0], [2.0]]).var(1)
ints = torch.tensor([[1, 2], [3, 4]])
ints.mean(1); ints.sum(1)
a.var(1); a.var_keepdim(1);
Tensor::new(&[[1f32], [2.]], &dev)?.var(1);
let ints = Tensor::new(&[[1i64, 2], [3, 4]], &dev)?;
ints.mean(1); ints.sum(1);
a.var(1) [2] torch.float32 [1.0, 1.0]
(2,1).var(1) [2] torch.float32 [nan, nan] [warning: var(): degrees of freedom is <= 0. Correction should be strictly less than the reduction factor (input numel divided by output numel). (Triggered internally at .../aten/src/ATen/native/ReduceOps.cpp:1927.)]
i64.mean(1) NotImplementedError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long
i64.sum(1) [2] torch.int64 [3, 7]
a.var(1) [2] F32 [1.0, 1.0]
a.var_keepdim(1) [2, 1] F32 [1.0, 1.0]
(2,1).var(1) [2] F32 [NaN, NaN]
i64.mean(1) [2] I64 [0.0, 0.0]
i64.sum(1) [2] I64 [3.0, 7.0]
var_keepdim 的注释是「Returns the unbiased variance」，实现是平方和除以 dim - 1，和 PyTorch 的默认 correction=1 相同；没有 correction 参数，要有偏方差就自己写 x.broadcast_sub(&x.mean_keepdim(d)?)?.sqr()?.mean_keepdim(d)?。只有一个元素时两边都是 NaN（0/0），PyTorch 多一条警告。
整数的 mean 两边差得远。PyTorch 拒绝；candle 的 mean 是 sum 之后乘上 1 / n，这一乘走 affine，系数 0.5 先转成 I64 变成 0，于是 CPU 上 [3, 7] 的均值是 [0, 0]，不报错。同一段代码在 Metal 上得到 [1, 3]（向零截断的 1.5 和 3.5），CPU 和 Metal 的结果不一样。这和 i64 * 0.5 是同一个机制，见整数运算；要平均值，先 to_dtype(DType::F32)。
下标与切片：i(..)
| PyTorch | candle | 结果形状 |
|---|---|---|
a[1] | a.i(1)? | [3] |
a[:, 2] | a.i((.., 2))? | [2] |
a[:, 1:] | a.i((.., 1..))? | [2, 2] |
a[:, :2]（含端点写法） | a.i((.., ..=1))? | [2, 2] |
a[1, 2] | a.i((1, 2))? | [] |
a[:, [2, 0]] | a.i((.., vec![2u32, 0]))? | [2, 2] |
a.narrow(1, 1, 2) | a.narrow(1, 1, 2)? | [2, 2] |
（每一行两边的值都相同，见实验输出。）IndexOp 的 i 接受一个下标或一个元组，元组的每一项被转成 TensorIndexer 的三种之一：usize 是 Select，Rust 的区间是 Narrow，Vec<u32>、&[u32]、&Tensor 是 IndexSelect。Tensor::index 依次把它们变成 narrow + squeeze、narrow、index_select，所以整数下标降秩、区间不降秩，这一点和 PyTorch 一致（burn 的整数下标不降秩，见 burn 卷的切片与索引）。i(..) 要先 use candle_core::IndexOp;。
PyTorch 有、candle 没有的写法：
a[:, 3]; a[:, 1:10]; a[:, 2:1]; a[:, -1]; a[:, ::2]; a[a >= 3]
a[:, torch.tensor([[0, 1]])]
a.i((.., 3)); a.i((.., 1..10)); a.i((.., 2..1)); a.i(2);
let idx2 = Tensor::new(&[[0u32, 1]], &dev)?;
a.i(&idx2);
let last = a.dim(D::Minus1)? - 1;
a.i((.., last));
a[:, 3] IndexError: index 3 is out of bounds for dimension 1 with size 3
a[:, 1:10] [2, 2] torch.float32 [1.0, 2.0, 4.0, 5.0]
a[:, 2:1] [2, 0] torch.float32 []
a[:, -1] [2] torch.float32 [2.0, 5.0]
a[:, ::2] [2, 2] torch.float32 [0.0, 2.0, 3.0, 5.0]
a[a >= 3] [3] torch.float32 [3.0, 4.0, 5.0]
a[:, torch.tensor([[0, 1]])] [2, 1, 2] torch.float32 [0.0, 1.0, 3.0, 4.0]
a.i((.., 3)) Err: narrow invalid args start + len > dim_len: [2, 3], dim: 1, start: 3, len:1
a.i((.., 1..10)) Err: narrow invalid args start + len > dim_len: [2, 3], dim: 1, start: 1, len:9
a.i((.., 2..1)) [2, 0] F32 []
a.i(2) Err: narrow invalid args start + len > dim_len: [2, 3], dim: 0, start: 2, len:1
a.i(&idx2d) Err: multi-dimensional tensor indexing is not supported
a.i((.., last)) [2] F32 [2.0, 5.0]
- 越界：单个下标两边都报错，PyTorch 抛
IndexError，candle 返回Err，报错是narrow的，因为整数下标就是narrow(dim, n, 1)。区间越界不一样：PyTorch 和 Python 的切片一样把1:10截到1:3，candle 报错。从 PyTorch 搬x[:, :max_len]这种「不够长就全取」的写法，要自己写..max_len.min(x.dim(1)?)。反向区间2..1两边都得到空张量（candle 用saturating_sub算长度）。 - 负下标：区间和下标都是
usize，a.i((.., -1))编译不过（the trait std::ops::Neg is not implemented for usize），要用a.dim(D::Minus1)? - 1算出来。 - 步长、布尔掩码、二维下标张量都没有。
a[:, ::2]用index_select拿一组算好的下标；a[a >= 3]的结果长度取决于数据，candle 没有对应的算子，通常改写成形状不变的where_cond（见下文），或者to_vec1到主机上筛。
index_select 与 embedding：下标的 dtype
a.index_select(1, torch.tensor([2, 0, 2]))
a.index_select(1, torch.tensor([2, 0], dtype=torch.int32))
a.index_select(1, torch.tensor([2, 0], dtype=torch.uint8))
a.index_select(1, torch.tensor([2.0, 0.0]))
a.index_select(1, torch.tensor([5]))
a.index_select(&Tensor::new(&[2u32, 0, 2], &dev)?, 1);
a.index_select(&Tensor::new(&[2i64, 0], &dev)?, 1);
a.index_select(&Tensor::new(&[2u8, 0], &dev)?, 1);
a.index_select(&Tensor::new(&[2i32, 0], &dev)?, 1);
a.index_select(&Tensor::new(&[2f32, 0.], &dev)?, 1);
a.index_select(&Tensor::new(&[5u32], &dev)?, 1);
a.index_select(1, [2,0,2]) [2, 3] torch.float32 [2.0, 0.0, 2.0, 5.0, 3.0, 5.0]
index_select int32 ids [2, 2] torch.float32 [2.0, 0.0, 5.0, 3.0]
index_select uint8 ids RuntimeError: index_select(): Expected dtype int32 or int64 for index
index_select float ids RuntimeError: index_select(): Expected dtype int32 or int64 for index
index_select oob 5 RuntimeError: INDICES element is out of DATA bounds, id=5 axis_dim=3
a.index_select(&ids, 1) [2, 3] F32 [2.0, 0.0, 2.0, 5.0, 3.0, 5.0]
a.index_select(i64 ids, 1) [2, 2] F32 [2.0, 0.0, 5.0, 3.0]
a.index_select(u8 ids, 1) [2, 2] F32 [2.0, 0.0, 5.0, 3.0]
a.index_select(i32 ids, 1) Err: unsupported dtype F32 for op index-select
a.index_select(f32 ids, 1) Err: unsupported dtype F32 for op index-select
a.index_select(oob 5, 1) Err: index-select invalid index 5 with dim size 3
参数顺序反过来了：PyTorch 是 index_select(dim, index)，candle 的 index_select 是 (indexes, dim)，下标张量必须是一维。两边接受的下标 dtype 也不同：PyTorch 是 int32 / int64，candle 是 U8 / U32 / I64，没有 I32。
I32 那一行的报错写的是 F32。CPU 后端的 index_select 按下标的 dtype 分派，落到最后一支时用的是数据的 dtype：
_ => Err(Error::UnsupportedDTypeForOp(self.dtype(), "index-select").bt()),
self 是被取值的张量，不是 ids，所以下标 dtype 不对时，报错里出现的是数据的 dtype（gather 同样）。这和创建与 dtype 里的整数字面量凑在一起，就是一个难查的错误：
let ids = Tensor::new(&[2, 0], &dev)?;
a.index_select(&ids, 1);
emb.embedding(&ids);
a.index_select(&ids, 1): Err("unsupported dtype F32 for op index-select")
emb.embedding(&ids): Err("unsupported dtype F32 for op index-select")
[2, 0] 是 I32，报错却让人去查 F32 的数据。看到 unsupported dtype <数据的 dtype> for op index-select，先看下标的 dtype；写 &[2u32, 0]。Metal 上同样的输入报 Metal contiguous index_select I32 F32 not implemented，两个 dtype 都列了出来。
越界时 CPU 返回 Err，Metal 不报错，把下标截到最后一个：reduce metal 的输出是 a.index_select(oob 5, 1) [2, 1] F32 [2.0, 5.0]，取的是第 2 列。Metal 内核 index 里写的是 min(input_ids[id_i], (INDEX_TYPENAME)(src_dim_size - 1))。
-1 不是最后一个
PyTorch 的高级索引接受负下标（emb[[-1, 2]] 取最后一行），index_select 不接受。candle 的下标是无符号或 I64：
emb = torch.arange(6, dtype=torch.float32).reshape(3, 2)
emb.index_select(0, torch.tensor([-1, 2])); emb[torch.tensor([-1, 2])]
let emb = Tensor::arange(0f32, 6., &dev)?.reshape((3, 2))?;
let ids_i64 = Tensor::new(&[-1i64, 2], &dev)?;
let ids_u32 = ids_i64.to_dtype(DType::U32)?;
emb.index_select(&ids_i64, 0); emb.index_select(&ids_u32, 0);
emb.index_select(&Tensor::new(&[-100i64, 2], &dev)?.to_dtype(DType::U32)?, 0);
torch index_select [-1,2]: IndexError index out of range in self
torch emb[[-1,2]]: [[4.0, 5.0], [4.0, 5.0]]
ids_u32 = [4294967295, 2]
i64 ids: Err("index-select invalid index 18446744073709551615 with dim size 3")
u32 ids: Ok([[0.0, 0.0], [4.0, 5.0]])
-100 as u32 ids: Err("index-select invalid index 4294967196 with dim size 3")
I64 的 -1 被当成 usize 解释成 2⁶⁴−1，报越界。转成 U32 之后（to_dtype 整数之间回绕，见创建与 dtype）变成 4294967295，即 u32::MAX，这时不报错，那一行变成全 0。-100（transformers 里 ignore_index 的常用值）转成 U32 不是最大值，照常报越界。全 0 来自 candle 有意设的一个规则：下标等于该类型的最大值时输出 0。
#2940「Indexing with max-value results in zero/no-op.」（首次进入 0.9.2-alpha.1，第一个正式版是 0.10.0）之前，CPU 的 IndexSelect、Gather 把每个下标 as_usize() 之后直接和维度长度比较，u32::MAX 和其他越界下标一样返回 InvalidIndex。这个 PR 在比较之前先判断 index == I::max_value()：index_select、gather 在那里写 0，scatter 跳过那个位置；CUDA 和 Metal 的内核一起改了（Metal 里就是上面 index 内核开头的 if (input_ids[id_i] == max_value<INDEX_TYPENAME>())）。PR 加的测试用 Tensor::new(&[u32::MAX, 2u32, u32::MAX], device) 做 index_select，断言结果是 [[0.0, 0.0], [4.0, 5.0], [0.0, 0.0]]。PyTorch 没有这样的哨兵值，从 int64 下标转成 U32 时，原来的 -1 会悄悄落到这条规则上。
gather：参数顺序和越界
g = torch.tensor([[2, 0], [1, 1]])
a.gather(1, g); a.gather(1, torch.tensor([[3], [0]])); a.gather(1, g.to(torch.uint32))
let g = Tensor::new(&[[2u32, 0], [1, 1]], &dev)?;
a.gather(&g, 1);
a.gather(&Tensor::new(&[[3u32], [0]], &dev)?, 1);
a.gather(&Tensor::new(&[2u32, 0], &dev)?, 1);
a.gather(1, g) [2, 2] torch.float32 [2.0, 0.0, 4.0, 4.0]
a.gather(1, g_oob) RuntimeError: index 3 is out of bounds for dimension 1 with size 3
a.gather(1, g_uint32) RuntimeError: gather(): Expected dtype int32/int64 for index
a.gather(&g, 1) [2, 2] F32 [2.0, 0.0, 4.0, 4.0]
a.gather(&g_oob, 1) Err: gather invalid index 3 with dim size 3
a.gather(&g_wrong_rank, 1) Err: shape mismatch in gather, lhs: [2, 3], rhs: [2]
语义相同（结果形状等于下标张量，out[i][j] = a[i][g[i][j]]），参数顺序同样是 (indexes, dim)。gather 要求下标和数据秩相同，其余维度不超过数据的长度。
越界在 Metal 上没有检查，也不截断。gather_oob 在两个设备上各跑两次：
Cpu gather [[0],[3]]: Err("gather invalid index 3 with dim size 3") gather [[u32::MAX],[1]]: Ok([[0.0], [4.0]])
Metal { gpu_id: 4294968354 } gather [[0],[3]]: Ok([[0.0], [0.0]]) gather [[u32::MAX],[1]]: Ok([[0.0], [4.0]])
reduce metal 里的 [[3], [0]] 得到 [3.0, 3.0]：第 0 行的下标 3 读到了第 1 行的第 0 个元素。Metal 的 gather 内核直接算 src_i = (left_rank_i * src_dim_size + input_i) * right_size + right_rank_i 去读，越界的下标落到相邻的行，最后一行越界则读过缓冲区末尾（这次读到 0，没有保证）。和 index_select 截到最后一行一样，Metal 上下标的正确性要调用方自己保证。
where_cond：顺序相同，不广播
cond = a >= 3
torch.where(cond, a, -1.0); torch.where(cond, a, torch.full((2, 3), -1.0))
torch.where(torch.tensor([True, False, True]), a, -1.0)
torch.where(cond, a, torch.full((2, 3), -1.0, dtype=torch.float64))
a.masked_fill(cond, float("-inf"))
let cond = a.ge(3f64)?;
let neg = Tensor::full(-1f32, (2, 3), &dev)?;
cond.where_cond(&a, &neg);
cond.where_cond(&a, &Tensor::new(-1f32, &dev)?);
cond.where_cond(&a, &Tensor::new(-1f32, &dev)?.broadcast_as((2, 3))?);
cond.to_dtype(DType::U32)?.where_cond(&a, &neg);
cond.to_dtype(DType::F32)?.where_cond(&a, &neg);
cond.where_cond(&a, &neg.to_dtype(DType::F64)?);
Tensor::new(&[1u8, 0, 1], &dev)?.where_cond(&a, &neg);
a >= 3 [2, 3] torch.bool [False, False, False, True, True, True]
torch.where(cond, a, -1.0) [2, 3] torch.float32 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
torch.where(cond, a, full(-1)) [2, 3] torch.float32 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
torch.where(rowcond, a, -1.0) [2, 3] torch.float32 [0.0, -1.0, 2.0, 3.0, -1.0, 5.0]
torch.where(cond, a, f64) [2, 3] torch.float64 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
a.masked_fill(cond, -inf) [2, 3] torch.float32 [0.0, 1.0, 2.0, -inf, -inf, -inf]
a.ge(3.0) [2, 3] U8 [0.0, 0.0, 0.0, 1.0, 1.0, 1.0]
cond.where_cond(&a, &neg) [2, 3] F32 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
cond.where_cond(&a, &scalar) Err: shape mismatch in where_cond, lhs: [2, 3], rhs: []
cond.where_cond(&a, &bcast) [2, 3] F32 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
u32 cond.where_cond [2, 3] F32 [-1.0, -1.0, -1.0, 3.0, 4.0, 5.0]
f32 cond.where_cond Err: unsupported dtype F32 for op where-cond
cond.where_cond(&a, &neg_f64) Err: dtype mismatch in where, lhs: F32, rhs: F64
rowcond.where_cond (no bcast) Err: shape mismatch in where_cond, lhs: [3], rhs: [2, 3]
torch.where(cond, x, y) 写成 cond.where_cond(&x, &y)：条件是接收者，真取第一个参数、假取第二个，顺序和 PyTorch 一样（burn 的 mask_where 不一样，照搬会把两半对调，见 burn 卷的坑 mask_where 照搬 torch.where）。差别都在「不自动做」上：
- 比较运算
ge、lt、eq返回 U8，不是 bool；CPU 上条件可以是 U8 / U32 / I16 / I32 / I64，不能是浮点。 where_cond开头两次same_shape_binary_op，三个张量形状必须完全相同，不广播；标量要先建成张量再broadcast_as（广播视图可以直接用）。- 两个分支 dtype 必须相同，不提升。
PyTorch 里常见的 scores.masked_fill(mask, float("-inf"))，在 candle 里写成「建一个 -inf 张量，broadcast_as 到 scores 的形状，mask.where_cond(&neg_inf, &scores)」。candle_transformers 的 llama.rs 等 23 个模型文件各自带一个这样的 masked_fill 小函数（llama.rs 里是 mask.where_cond(&on_true, on_false)，on_true 是 broadcast_as 过的标量），注意力里怎么用它见源码卷的 quantized_llama。
对得上的数字
64×1000 的 f32 矩阵，元素是 sin(k)，输入只在 PyTorch 侧生成、写进 JSON，candle 侧读同一份（和 burn 卷归约与广播用的是同一个输入）：
out = {
"sum_dim1": x.sum(dim=1).tolist(), "mean_dim0": x.mean(dim=0).tolist(),
"sum": x.sum().item(), "mean": x.mean().item(), "var_dim1": x.var(dim=1).tolist(),
"max_dim1": x.max(dim=1).values.tolist(), "argmax_dim1": x.argmax(dim=1).tolist(),
}
let out = serde_json::json!({
"sum_dim1": x.sum(1)?.to_vec1::<f32>()?,
"mean_dim0": x.mean(0)?.to_vec1::<f32>()?,
"sum": x.sum_all()?.to_scalar::<f32>()?,
"mean": x.mean_all()?.to_scalar::<f32>()?,
"var_dim1": x.var(1)?.to_vec1::<f32>()?,
"max_dim1": x.max(1)?.to_vec1::<f32>()?,
"argmax_dim1": x.argmax(1)?.to_vec1::<u32>()?,
});
CPU（开不开 accelerate 结果相同），compare.py 容差 1e-5：
sum_dim1 n= 64 max|diff|=3.20e-06
mean_dim0 n=1000 max|diff|=1.68e-08
sum n= 1 max|diff|=5.75e-06
mean n= 1 max|diff|=9.00e-11
var_dim1 n= 64 max|diff|=1.19e-07
max_dim1 n= 64 max|diff|=0.00e+00
argmax_dim1 n= 64 max|diff|=0.00e+00
worst=5.75e-06 tol=1e-05 OK
max 和 argmax 只是比较，差 0。求和不是 0，因为 f32 加法不满足结合律，两边把 1000 个数加起来的顺序不同。拿同一份 f32 输入在 f64 里算参考，两边离它都有误差：PyTorch 最大 5.41e-6，candle 最大 5.03e-6。误差的量级看的是参与求和的数的绝对值之和：每行 1000 个 |sin(k)| 之和约 637，乘上 f32 的 ε（1.19e-7）约 7.6e-5，每行和的差 3.20e-6 远在其内；64000 个数的总和本身只有 0.374，所以 5.75e-6 的绝对误差看起来大。容差 1e-5 和 burn 卷那一页相同，理由也相同。
Metal 上同一份输入，sum 差 3.30e-5，超过 1e-5（MISMATCH），离 f64 参考 2.76e-5；其余各项不超过 3.81e-6。Metal 的求和顺序和 CPU 不同；64000 个数的绝对值之和约 4.07e4，乘上 ε 是 4.9e-3，3.30e-5 仍在 f32 的舍入范围内。在 Metal 上比较全局和，容差要按这个量级放宽。
动脑筋答案
nan.max(1) [2] F32 [2.0, NaN]
nan.argmax(1) [2] U32 [2.0, 0.0]
第一行 [1, NaN, 2] 的最大值是 2，下标 2，NaN 被跳过；第二行 [NaN, 0, 5] 的最大值是 NaN，下标 0。结果取决于 NaN 在不在第一个位置。CPU 后端的 ReduceIndex::fold_impl 从 val = src[0] 开始，遇到 f(val, s) 为真才替换，max 的 f 是 |x, y| x < y。val < NaN 和 NaN < y 都是假：NaN 不在开头时永远换不进来，在开头时永远换不出去。
Metal 上同一段代码又是另一个结果：
nan.max(1) [2] F32 [2.0, 5.0]
nan.argmax(1) [2] U32 [2.0, 2.0]
两行的 NaN 都被跳过。PyTorch 在两种情况下都返回 NaN（argmax 指向 NaN 的位置）。logits 里出现 NaN 时，PyTorch 的贪心解码会选到 NaN 所在的那个 token，candle 可能选别的 token，CPU 和 Metal 选的还可能不一样；要发现 NaN，得单独检查（比如 x.ne(&x)? 再 sum_all），不能指望 max 把它带出来。
小结
sum、mean、max、argmax和 PyTorch 一样去掉被归约的维度，*_keepdim对应keepdim=True；全部归约是sum_all/mean_all，得到 0 维张量。max只给值，argmax另调，下标是 U32；并列取第一个，两边相同。- NaN：PyTorch 的
max/argmax传播 NaN；candle CPU 的结果取决于 NaN 是否在第一个位置，Metal 跳过 NaN。 var无偏，没有correction；整数的mean不报错，CPU 上[3, 7]的均值是[0, 0]，Metal 上是[1, 3]。i(..)：整数降秩、区间不降秩；区间越界报错而不是截断；没有负下标、步长、布尔掩码、二维下标张量。index_select(ids, dim)、gather(ids, dim)参数顺序和 PyTorch 相反；下标 U8 / U32 / I64，没有 I32，dtype 不对时报错写的是数据的 dtype；u32::MAX下标输出 0；Metal 上越界不报错（index_select截断，gather读到别处）。cond.where_cond(&x, &y)的顺序和torch.where相同，但三者形状、两个分支的 dtype 都必须一致。- 64×1000 的归约在 CPU 上和 PyTorch 最大差 5.75e-6（容差 1e-5），来自求和顺序；
max/argmax差 0。