Turing.jl 项目中的自动微分测试重构实践

引言

在Julia生态系统的概率编程语言Turing.jl中，自动微分（AD）是核心功能之一。近期开发团队针对项目中的AD测试进行了深入讨论和重构，旨在优化测试效率并提高代码质量。本文将详细介绍这一技术改进的背景、思路和具体实施方案。

问题背景

Turing.jl项目中存在大量通过完整运行采样器来测试AD功能的测试用例。这种方式虽然全面，但存在明显的性能问题：

1. 计算资源消耗大：每次测试都需要运行完整的MCMC采样链
2. 测试时间长：影响开发迭代速度
3. 测试关注点混杂：将AD正确性测试与采样器正确性测试耦合在一起

技术改进方案

经过团队讨论，确定了以下改进方向：

分离测试关注点

将AD功能测试与采样器正确性测试分离：

1. 采样器正确性测试：只需针对一个基准AD后端（如ForwardDiff）进行完整采样测试
2. AD功能测试：通过直接验证梯度计算来测试各AD后端

建立标准化测试集

创建统一的测试用例集合，明确Turing.jl对AD系统的需求：

1. 定义一组标准模型和对应的LogDensityFunction
2. 这些测试用例将成为AD包集成测试的参考标准
3. 保持测试集的持续更新，反映Turing.jl的最新需求

技术实现细节

测试结构重构

1. 减少完整采样测试：仅对基准AD后端保留完整采样链测试
2. 增加直接梯度验证：使用logdensity_and_gradient函数直接测试各AD后端
3. 引入随机种子控制：使用StableRNGs确保测试可重复性

与DifferentiationInterface的整合

考虑与Julia生态中的DifferentiationInterface测试框架对接：

1. 使测试用例能被DifferentiationInterface的测试功能消费
2. 为未来全面采用DifferentiationInterface做准备
3. 当前仍需直接测试logdensity_and_gradient，因为部分后端尚未完全迁移

实施进展

目前已在实验性包ModelTests.jl中开始实现：

1. 收集典型模型作为测试用例
2. 设计兼顾当前需求和未来扩展的测试接口
3. 探索与DynamicPPL.TestUtils的整合方案

技术挑战与解决方案

后端兼容性问题

不同AD后端在Turing.jl中的集成方式各异：

1. ForwardDiff等传统后端使用专用扩展
2. Mooncake等新后端通过DifferentiationInterface集成
3. 需要设计兼容两种方式的测试框架

测试准确性保障

1. 对基准AD后端保留完整采样测试
2. 梯度测试采用数值验证方法
3. 建立结果一致性检查机制

未来工作方向

1. 推动LogDensityProblemsAD全面采用DifferentiationInterface
2. 完善测试用例集合，覆盖更多边界情况
3. 优化测试执行效率，缩短CI时间
4. 建立AD性能基准测试体系

总结

Turing.jl项目的AD测试重构工作体现了软件工程中"关注点分离"的重要原则。通过将AD功能测试与采样器测试解耦，不仅提高了测试效率，还使项目对AD后端的质量要求更加明确。这一改进将为Turing.jl的长期维护和生态扩展奠定坚实基础，同时也为Julia生态中其他需要AD支持的项目提供了有价值的参考模式。