架构概览

TensorPlay 采用严格解耦的分层架构设计。它由四个核心库组成，具有清晰的边界和单向依赖关系。

4 个核心库

1. P10（张量计算引擎）

• 角色：基础的“计算引擎”。
• 设计理念：
  • 硬件抽象：使用 Tensor 接口和 TensorImpl 多态性来支持多种硬件后端（CPU、CUDA、自定义边缘芯片），而无需更改前端 API。
  • 零微分逻辑：专注于高效的张量内核和内存管理，作为所有其他层的稳定基石。
  • 调度器模式：将算子定义与特定于设备的实现解耦，允许轻松集成 MKL 或 cuDNN 等库。

2. TPX（自动微分引擎）

• 角色：“微分层”。
• 设计理念：
  • 解耦的自动微分：作为轻量级扩展层实现，而不是烘焙到张量核心中。它仅在 requires_grad=True 时跟踪操作。
  • 显式图构建：专为教育清晰度而设计，允许用户检查计算图是如何构建的以及在反向传播期间是如何遍历的。
  • 可插拔引擎：可以替换或扩展不同的微分模式（例如高阶导数），而不影响底层的 P10 引擎。

3. Stax（静态图加速器）

• 角色：“优化层”。
• 设计理念：
  • 优化优先：专注于静态图捕获、算子融合和即时 (JIT) 编译，以最大限度地减少 Python 开销。
  • 独立路径：在与 TPX/NN 分离的依赖链上运行，使其成为可以根据性能需求添加或删除的模块化组件。
  • 编译器集成：设计用于在未来对接 MLIR 或 TVM 等高级编译器后端。

4. NN（神经网络库）

• 角色：“业务层”。
• 设计理念：
  • 用户友好的抽象：为 Linear、Conv2d 和 Optimizers 等高级组件提供熟悉的、兼容 PyTorch 的接口。
  • 蓝图方法：每一层都被设计成清晰、可读的蓝图，演示复杂的神经网络组件是如何从基本的张量操作构建的。
  • 纯依赖：严格依赖 P10 和 TPX 的公共 API，确保它保持为高级建模的可选、非侵入式层。

---

依赖图

graph TD
    NN[NN: Neural Networks] --> TPX[TPX: Autograd]
    NN --> P10[P10: Computation]
    TPX --> P10
    Stax[Stax: Static Graph] --> P10
    
    style P10 fill:#f9f,stroke:#333,stroke-width:2px
    style TPX fill:#bbf,stroke:#333,stroke-width:1px
    style Stax fill:#bfb,stroke:#333,stroke-width:1px
    style NN fill:#fbb,stroke:#333,stroke-width:1px

为什么采用这种架构？

1. 解耦：每个库都可以独立开发、测试和优化。
2. 可扩展性：添加新的硬件后端仅需更改 P10。添加新的微分模式仅影响 TPX。
3. 性能：用户只需为他们使用的功能付费。纯计算任务不会加载自动微分引擎或静态图编译器。
4. 教育价值：通过分离这些关注点，TensorPlay 使开发人员能够轻松理解现代深度学习框架的内部结构。