Abstract
The standard dual-numbers construction works well for forward-mode automatic differentiation (AD) and is attractive due to its simplicity; recently, it also has been adapted to reverse-mode AD, but practical performance, especially on array programs, leaves a lot to be desired. In this paper we introduce first-class support for multidimensional arrays in dual-numbers reverse-mode AD with little to no performance overhead. The algorithm consists of three loosely-coupled components: a semantics-preserving vectorisation code transformation (the bulk-operation transform or BOT), a fairly straightforward lifting of the basic dual-numbers reverse AD algorithm to a mostly first-order array language, and symbolic interpretation to achieve an end-to-end compilation pipeline. Unfortunately, we lose some of the nice generalisable aspects of dual-numbers AD in the process, most importantly support for higher-order code.
We do support some higher-order array combinators, but only a carefully-chosen set: ‘build’ (elementwise array construction), ‘gather’ and ‘scatter’. In return, the BOT can eliminate the essential (for AD) higher-orderness of the input program, meaning that AD gets essentially presented with a first-order program. This allows the naive trick of lifting dual numbers to “dual arrays” to work without much modification.
Archiv URL: https://arxiv.org/abs/2507.12640.
We have also published (a first version of) an accompanying Haskell library that supports array AD,
called horde-ad:
https://hackage.haskell.org/package/horde-ad.
It comes with an accompanying highly-performant
shapely multdimensional array library called ox-arrays:
https://hackage.haskell.org/package/ox-arrays.
This technical report documents our progress in applying the dual-numbers construction to reverse AD for array programs. While there is material for more than a single paper here, we are still unsatisfied with the elegance of the algorithm design, and we do not yet know if these algorithms can be improved (without extensive special-casing) to reach performance competitive with popular machine learning frameworks. If you would like to collaborate with us in improving the algorithms or the presentation as one or more papers, please reach out! (Email addresses are in the paper.)