MLX项目中AddMM操作的JVP实现问题分析

背景介绍

MLX是苹果推出的一个深度学习框架，它提供了自动微分功能，这是现代深度学习框架的核心组件之一。自动微分主要分为两种模式：前向模式（Forward Mode）和反向模式（Reverse Mode）。在MLX框架中，这两种模式分别通过jvp（Jacobian-vector product）和vjp（vector-Jacobian product）函数实现。

问题现象

在使用MLX框架时，开发者发现一个有趣的现象：对于包含线性层（nn.Linear）的简单神经网络，使用vjp函数计算向量-雅可比积可以正常工作，但使用jvp函数计算雅可比-向量积时却会抛出"Not implemented for AddMM"的错误。

技术分析

这个问题的根源在于MLX框架中AddMM操作（矩阵乘法加偏置的融合操作）的前向模式自动微分实现尚未完成。AddMM是深度学习中非常基础且重要的操作，它代表了以下计算：

output = input @ weight.T + bias

在反向模式自动微分（vjp）中，MLX已经完整实现了AddMM操作的梯度计算。但在前向模式自动微分（jvp）中，该操作的实现尚未完成，导致框架在尝试计算时会抛出未实现的错误。

影响范围

这个问题会影响所有需要使用前向模式自动微分的场景，包括：

1. 高阶导数计算
2. 某些优化算法中需要同时使用前向和反向模式的情况
3. 需要计算完整Hessian矩阵的应用

解决方案

根据MLX开发团队的反馈，这个问题将会很快得到修复。AddMM操作的JVP实现相对直接，可以预期在不久的将来就会包含在框架更新中。

临时替代方案

在官方修复发布前，开发者可以考虑以下替代方案：

1. 使用反向模式自动微分（vjp）替代
2. 手动实现线性层的前向模式微分
3. 分解AddMM操作为单独的矩阵乘法和加法操作

总结

MLX作为一个新兴的深度学习框架，正在逐步完善其功能。这个AddMM的JVP实现问题反映了框架开发过程中的常见情况 - 优先实现最常用的功能（如反向模式自动微分），然后再完善其他功能。随着框架的成熟，这类问题将会越来越少。

对于开发者而言，理解自动微分的两种模式及其实现原理，有助于更好地使用深度学习框架，并在遇到类似问题时能够快速定位原因和找到解决方案。