结论：tinygrad 的 nn/__init__.py 是通用模型层库，它的层按「形状契约 + 状态契约」工作：BatchNorm 的 running_mean/running_var/num_batches_tracked 是推理结果相关状态，Conv1d/ConvTranspose1d 复用二维实现，Embedding 默认路径对越界索引返回零向量而非显式拒绝。openpilot 是通过 ONNX 构建已捕获模型，不能把层库的训练能力直接等同于车端在线学习。
组件与关键函数
基线是 source/ 固定快照，逐类解析公开层和两个 Embedding 内部路径。各层的形状契约必须写清楚，不能只按名字替换。
| 组件 | 实现与输入输出契约 |
|---|---|
BatchNorm |
通道轴固定为 1，均值/方差跨其余轴。训练时以 momentum 更新累计统计量；方差更新含 x.numel()/(x.numel()-x.shape[1]) 无偏校正，单通道单样本等分母为零的输入不属于可直接保证的有效域。BatchNorm2d/3d 只是别名。 |
Conv1d / Conv2d |
一维接口以单元素 kernel tuple 复用通用卷积层。权重为 (out_channels, in_channels/groups, *kernel)；padding='same' 只接受 stride=1，按反向空间轴生成不对称 pad。 |
ConvTranspose1d / ConvTranspose2d |
初始化后重建权重为 (in_channels, out_channels/groups, *kernel)，调用 conv_transpose2d() 并传入 stride、padding、dilation、groups 与 output_padding。加载普通卷积权重不能只凭同名字段互换。 |
Linear |
初始化权重 (out_features, in_features)，前向转置后调用 x.linear()，在最后一维做变换；支持前置批/时间维。 |
LSTMCell |
weight_ih/weight_hh 产生四门 i/f/g/o，输入为 (batch,input_size)，可选 (h,c) 形状为 (batch,hidden_size)；省略时创建零状态，调用者负责跨时间步保存与回传状态。 |
归一化层各有约定：GroupNorm 把输入 reshape 为 (N, groups, -1)、InstanceNorm 为 (N, channels, -1)，借 layernorm() 沿末维归一化再还原形状；输入通道数/分组整除与 NCHW 布局是适配契约。LayerNorm 断言输入尾维与 normalized_shape 完全一致，LayerNorm2d 把 NCHW 暂转 NHWC 归一化通道再转回，二者都不更新运行统计量。RMSNorm 按末维平方均值计算逆 RMS，先转 float 计算、再转回输入 dtype，半精度目标后端需单独比对数值误差。
形状契约上还有几处必须显式记住：BatchNorm2d/3d 只是同一类的别名；ConvTranspose 系列在重建权重后调用 conv_transpose2d()，并把 output_padding 一并传入，普通卷积权重不能只凭字段同名就替换；Linear 前向会先转置权重、在最后一维做变换，因此前置的批/时间维是允许的。
Embedding 的两条内部路径
Embedding.__call__() 仅接受整数索引。默认通过与 arange(vocab) 全量比较、选取权重再沿词表归约，输出形状为 indices.shape + (embed_size,)；索引超出词表时默认路径得到零向量，并非显式拒绝。当 USE_ATOMICS=1 时把前向包成 CALL，反向先清零 FP32 梯度，再按 token 原子 scatter-add。
这里有一条语义差异：原子实现仅列出 CPU/NULL/AMD，词表分片只支持轴 0、其他分片会断言；索引在反向内 clip 到词表边界，因此无效 token 在默认前向和原子反向可能产生不同语义。插件入口应先验证索引在 0 到 vocab_size 范围内。
状态管理方面，LSTMCell 在省略初始状态时创建零状态，返回新的 (h,c)，跨时间步的保存与回传由调用者负责；LayerNorm、LayerNorm2d、GroupNorm、InstanceNorm 与 RMSNorm 都不更新运行统计量，只有 BatchNorm 会写入累计统计量，这也是它在导出时必须与权重一起版本化的原因。
状态、硬件与插件契约
BatchNorm 的运行统计量属于推理结果相关状态，模型导出必须与权重、eps、仿射配置一起版本化；切换 TRAINING 会改变统计来源及写入行为。TRAINING 与 USE_ATOMICS 默认都关闭。Conv2d、Linear 等初始化使用随机权重，加载模型工件后才能代表实际训练模型。
| 契约项 | 必须声明的内容 |
|---|---|
| 张量 ABI | 输入 shape/dtype/布局、权重 hash、运算版本、目标后端、数值容差。 |
| 状态 | 运行统计量、状态复位条件、最坏时延。 |
| 权限 | 模型层实现不得获得 CAN 发送、panda 安全配置或车辆控制发布权限。 |
| 原子路径 | 若启用原子自定义内核，需单独审计设备支持、重复索引并发与越界 token 行为。 |
「万物即插件」宜以可验证的模型工件和张量 ABI为边界，而不是把每个层都做成车端可热替换插件。相机/推理硬件插件提供帧、时间戳、缓冲所有权和执行能力；模型层实现只消费签名/版本绑定的工件。
反过来说，层库能训练、能加载随机权重，并不代表它在车端可用；只有把权重、运行统计量与运算版本一起冻结、并对目标后端做数值与时延验收之后，这些层才具备进入推理链的前提。
与 openpilot 的关系与验证范围
compile_model() 构建 ONNX/JIT 工件；运行期 ModelState.run() 调用已捕获模型。对 openpilot 与 examples 目录的 nn.<层名> 搜索未发现直接实例化上述层；这只说明可见 Python 源码，不能推断缺失的 LFS 模型图里没有对应 ONNX 算子。
验证层面，以 PYTHONPATH=. DEV=PYTHON RUN_SLOW=1 运行六项上游定向测试：BatchNorm2d、Linear、LayerNorm2d、RMSNorm、Embedding、LSTMCell，全部通过（51.793 秒）。转置卷积、原子路径、分片、多后端、真实模型数值与目标车载硬件仍未验证；LFS 模型实体缺失，不能声称整模型或整车验证。
常见问题
Embedding 遇到越界索引会发生什么？
默认路径会得到零向量，并非显式拒绝：_embedding_fwd() 通过与 arange(vocab) 全量比较、选取权重再沿词表归约。但索引在原子反向内会被 clip 到词表边界，所以无效 token 在默认前向和原子反向可能产生不同语义，插件入口应先验证索引在 0 到 vocab_size 范围内。
BatchNorm 的运行统计量需要一起导出吗？
需要。running_mean/running_var/num_batches_tracked 是推理结果相关状态，源码要求模型导出必须与权重、eps、仿射配置一起版本化；切换 TRAINING 会改变统计来源及写入行为。TRAINING 默认关闭。
Conv1d 能直接用 Conv2d 的权重换吗？
不能只凭同名字段互换。Conv1d 以单元素 kernel tuple 复用通用卷积层；而 ConvTranspose1d/2d 会在初始化后重建权重为 (in_channels, out_channels/groups, *kernel) 并调用 conv_transpose2d()，加载普通卷积权重不能只靠字段名相同就替换。
可以把每个 nn 层做成车端热替换插件吗？
不建议。源码建议以可验证的模型工件和张量 ABI 为边界，而不是把每个层都做成车端可热替换插件；契约至少包括输入 shape/dtype/布局、权重 hash、运算版本、目标后端、数值容差、运行统计量、状态复位条件和最坏时延。模型层实现不得获得 CAN 发送、panda 安全配置或车辆控制发布权限。
延伸阅读
- tinygrad 的 PrivateUse1/ATen 别名与设备 —— 承接层实现与 Torch 前端桥接的边界。
- tinygrad dtype 的 UOp 降解与模拟 —— 解释 RMSNorm 等为何要转 float 计算再回转。
- 端侧推理工具链快速上手 —— 说明层库之外真正进车端的模型工件与 ABI。
- tinygrad JIT 捕获与回放 —— 对照构建期 ONNX 工件与运行期捕获模型。
参考资料
- 项目 02 ——
openpilot-research/TINYGRAD_NN_LAYERS_STATE_SHAPE_EMBEDDING_CONTRACT_DEEP_DIVE.md（取用了各层形状/状态契约、Embedding 两条内部路径、插件边界与六项上游测试结论）。 - 项目 02 ——
openpilot-research/TINYGRAD_DTYPE_UOP_LOWERING_EMULATION_DEEP_DIVE.md（仅用于确认 dtype 降解专题的接续关系，未取用其未公开内容）。
说明：本文为对固定源码快照的静态结构分析，未运行完整测试套件、未连接硬件、未验证目标车型的真实参数组合。文中所有阈值、行号索引与配置均出自源文档，标注为未验证的结论保留其原有边界。
更新记录
首次发布 2026-10-11 最近核对：2026-10-11
发表回复