PyTorch核心张量算子

覆盖：扩散模型 / DiT / 图像视频 / 机器人动作序列

约定记号：
$\mathbf{X},\mathbf{Y}$ 代表张量；$N,B,C,T,H,W$ = 批次、通道、时序、高、宽

一、维度变换算子（特征排布，最高频）

1. tensor.view(*shape) / torch.reshape

数学：重塑张量元素排布，元素顺序不变
设原张量 $\mathbf{X}\in {\mathbb{R}}^{a\times b\times c}$
$$\mathrm{reshape}(\mathbf{X},[a\cdot b,c])\in {\mathbb{R}}^{ab\times c}$$
限制：元素总数不变；view要求内存连续，reshape自动处理连续性。
场景：DiT时空维度打平、latent变形、动作序列维度重组。

2. torch.permute(d₀,d₁,…)

维度重排列，无复制（非连续张量）
例：$\mathbf{X}\in {\mathbb{R}}^{B,C,H,W}$
$$\mathrm{permute}(\mathbf{X},[B,H,W,C])$$
场景：图像 BCHW ↔ BHWC、视频调换时序/空间维度、注意力维度对齐。

3. torch.unsqueeze(dim) / squeeze(dim)

$$\mathrm{unsqueeze}(\mathbf{X},d):{\mathbb{R}}^{D_0\times \dots D_k}\to {\mathbb{R}}^{D_0\times \cdots \times 1\times \dots D_k}$$
新增长度为1的维度；squeeze移除长度=1维度。
场景：单张图像扩batch维、给latent增加时序维度适配视频模型输入。

4. torch.stack([X₁,X₂,…,Xₙ], dim)

新增维度堆叠
${\mathbf{X}}_i\in {\mathbb{R}}^{H,W}$
$$\mathbf{Z}=\mathrm{stack}([{\mathbf{X}}_1,{\mathbf{X}}_2],0)⟹\mathbf{Z}\in {\mathbb{R}}^{2\times H\times W}$$
公式表达：
$$\mathbf{Z}[0]={\mathbf{X}}_1,\mathbf{Z}[1]={\mathbf{X}}_2$$
👉 KV Cache组装、多帧latent构建时序维度。

5. torch.cat([X₁,X₂], dim)

现有维度拼接，不新增维度
${\mathbf{X}}_1,{\mathbf{X}}_2\in {\mathbb{R}}^{B,C,H,W}$
$$\mathbf{Z}=\mathrm{cat}([{\mathbf{X}}_1,{\mathbf{X}}_2],2)\in {\mathbb{R}}^{B,C,2H,W}$$
场景：拼接条件嵌入、前后帧特征、历史动作序列。

二、采样 & 扩散核心数值算子

1. torch.randn_like(X)

$$\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})$$
标准高斯噪声，扩散模型根基：
$${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{\epsilon }$$

2. torch.lerp(a, b, weight) 线性插值

$$\mathrm{lerp}(\mathbf{a},\mathbf{b},w)=(1-w)\cdot \mathbf{a}+w\cdot \mathbf{b}$$
应用：扩散采样更新latent、流匹配积分、动作轨迹平滑、帧插值。

3. torch.clamp(x, min, max)

$$\mathrm{clamp}(x,a,b)=\{\begin{array}{ll}a& x<a\\ x& a\le x\le b\\ b& x>b\end{array}$$
场景：VAE解码值域限制防止数值溢出黑图、机器人动作限幅。

4. $\text{torch.sqrt},\exp ,\log$

扩散调度器通用：
$${\sigma }_t=\sqrt{1-{\bar{\alpha }}_t},{\bar{\alpha }}_t=\prod _{s=1}^t{\alpha }_s$$
Flow Matching / Cosmos世界模型大量使用对数、指数计算SNR、时间条件。

三、掩码与条件选择

1. torch.masked_fill(mask, value)

$${\mathbf{Y}}_{i,j}=\{\begin{array}{ll}v& {\mathrm{mask}}_{i,j}=True\\ {\mathbf{X}}_{i,j}& {\mathrm{mask}}_{i,j}=False\end{array}$$
场景：Inpaint区域填充噪声；注意力padding位置填充 $-\infty$。

2. torch.where(cond, x, y)

$$\mathrm{where}(c,\mathbf{x},\mathbf{y}{)}_i=\{\begin{array}{ll}{\mathbf{x}}_i& c_i=\mathrm{True}\\ {\mathbf{y}}_i& c_i=\mathrm{False}\end{array}$$
场景：根据掩码混合原图latent与噪声latent。

四、注意力矩阵运算（DiT / Transformer / KV Cache）

1. torch.matmul(a,b) / @

矩阵乘法
$$\mathbf{C}=\mathbf{A}\mathbf{B},C_{ij}=\sum _k{A}_{ik}{B}_{kj}$$
自注意力核心：
$$\mathrm{Attn}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d_h}}\right)\mathbf{V}$$

2. torch.transpose(d1,d2)

$$({\mathbf{X}}^{\top }{)}_{ij}=X_{ji}$$
常用写法：K.transpose(-1,-2) 适配QK相乘。

3. F.softmax(x, dim)

$$\mathrm{softmax}(\mathbf{x}{)}_i=\frac{e^{x_i}}{{\sum }_j{e}^{x_j}}$$

五、重采样算子（图像/视频latent缩放）

F.interpolate

连续空间重采样，以bilinear为例：
$$\hat{f}(u,v)=\sum _i\sum _j{w}_{ij}\cdot f(x_i,y_j)$$
场景：VAE上下采样、latent分辨率调整、视频时序缩放。

六、时序序列算子（机器人动作 / 视频）

1. torch.diff(X, dim)

相邻差分
$$\Delta {\mathbf{x}}_t={\mathbf{x}}_{t+1}-{\mathbf{x}}_t$$
用途：由位置序列求解速度、加速度。

2. torch.cumsum(X, dim)

累积和
$$\mathrm{cumsum}(\mathbf{X}{)}_t=\sum _{k=1}^t{\mathbf{X}}_k$$
Flow Matching积分、由速度积分得到轨迹。

七、归一化变换（图像VAE预处理）

x.sub_(mean).div_(std)

$$\hat{x}=\frac{x-\mu }{\sigma }$$
反归一化：
$$x=\hat{x}\cdot \sigma +\mu$$
常用：图像 $[0,1]\leftrightarrow [-1,1]$ 转换输入VAE。

八、内存连续性

tensor.contiguous()

张量底层内存由非连续 → 连续排布。
数学上不改变张量数值，改变存储布局；
permute/transpose后直接view会报错，必须调用。

熟练运用PyTorch张量算子完成扩散模型、视频时序与机器人动作序列开发：
通过stack/cat/permute/view/unsqueeze实现张量维度重组；依托$randn\_like$生成高斯噪声，结合$\mathrm{lerp}$线性插值实现扩散与流匹配采样迭代；利用$\mathrm{clamp}$控制数值范围，$masked\_fill/where$实现条件生成与图像修复；基于$matmul/transpose/softmax$搭建自注意力计算；使用$\mathrm{interpolate}$完成潜空间分辨率重采样；借助$diff/cumsum$对动作序列做差分与积分运算；同时掌握contiguous()处理维度置换后的内存连续性问题。