Rust in Everything
基于 candle 0.11.0

张量

这一部分走一遍张量主线:建张量、算、读回主机。candle 的 Tensor 和 torch.Tensor 一样把形状、dtype、设备都放在运行期,大多数操作能一一对上;不同的地方集中在「PyTorch 替你做了、candle 不做」的那几件事上:dtype 从 Rust 类型来、不提升,运算符不广播,切片不截断,整数不转浮点,设备之间不搬。每页都是同一个任务的两份程序,数字从实测输出里贴。candle 内部怎么实现张量,见源码卷的张量。

  1. 创建与 dtype:Tensor::new 的 dtype 由 Rust 类型定,1.0 是 F64、1 是 I32;没有默认 dtype、没有类型提升,和标量运算只认 f64;to_dtype 浮点转整数饱和;f32 转 BF16 / F16 的舍入两边逐位相同。
  2. 形状与广播:+ 不广播,要写 broadcast_add;matmul 不接一维、不广播批维,expand 之后直接乘在 gemm 后端上静默算错;reshape 的 -1 写成 ();t() 转置最后两维。
  3. 归约与索引:sum、argmax 去掉维度和 PyTorch 一样,argmax 是 U32;i(..) 区间越界报错不截断;index_select 的下标没有 I32、报错写错 dtype;where_cond 不广播;NaN 的 max 依位置而定。
  4. 整数运算:整数除法向零截断,没有 // 和 %;整数张量乘浮点标量在 CPU 和 Metal 上结果不同;溢出与回绕。
  5. 设备:Device::Cpu 和 new_metal(0),两次 new_metal(0) 是两个设备;不做隐式搬运;to_device;Metal 读回才等计算完成。
flowchart TB
  data["Rust 数据"] -->|"Tensor new"| t["Tensor"]
  dev["Device"] -.-> t
  t -->|"同 dtype 同形状"| op["运算 归约 索引"]
  op --> r["结果 Tensor"]
  r -->|"to_vec1 T"| host["主机 Vec"]
  t -.->|"dtype 不同"| e1["dtype mismatch"]
  t -.->|"形状不同"| e2["shape mismatch"]
  t -.->|"设备不同"| e3["device mismatch"]

💬 关联讨论

发起讨论
加载中...