Differentiable Fortran with LFortran and Enzyme
4 hours ago
- #automatic-differentiation
- #fortran
- #scientific-computing
- 可微分的Fortran允许通过现有的Fortran、C或C++模拟代码使用LFortran和Enzyme实现反向传播，从而能与JAX和PyTorch等现代机器学习框架集成。
- Enzyme在LLVM IR层面应用自动微分，可以对任何编译到LLVM的代码进行微分，避免了重写经过验证的遗留物理代码的需求。
- 使用LFortran而非Flang，是因为它生成的LLVM IR更简洁，使用纯指针运算，使Enzyme更容易追踪内存访问，而不涉及不透明的运行时调用。
- 编译流水线涉及六个步骤：将Fortran转换为LLVM IR、轻度优化、与C包装器链接、通过Enzyme合成导数、最终优化为共享库。
- Enzyme计算的梯度与解析导数匹配，精度很高（相对误差约6e-12），优于有限差分法，后者在截断误差与舍入误差之间存在权衡问题。
- Tesseract将微分求解器封装为自定义JAX原语，允许Fortran代码在JAX工作流中作为可微分层使用，这在一个热取证反问题中得到了演示。
- 一个案例研究展示了如何使用L-BFGS-B优化从100个含噪声观测中重建900个元素的初始温度场，通过反向模式自动微分计算梯度，速度比有限差分法提升约450倍。
- 该方法仍处于实验阶段，需要调整代码以避免不可追踪的结构（如可分配数组），在自动微分前进行轻度优化，并在IR层面进行调试，但它提供了与其他自动微分源的组合性。