Feasible to perform pullback preparation at compile time?

[MasonProtter]

Are there plans / ambitions to have a compile-time preparation step for Mooncake.jl? Or is it heavily baked into the design that you need values to construct the tape?

[willtebbutt]

Hi @MasonProtter . There are two kinds of preparation which happen:

1. rule derivation / code generation, and
2. allocation of (co)tangent / shadow memory in which to store results.

The former is entirely a function of the types of the arguments, so can happen once those are available. The tangent memory is necessarily allocate at runtime (as in Enzyme) because you need e.g. size information about arrays in order to allocate it.

At the minute, in both DI and Mooncake's own interfaces, these two steps get bundled together.

Additionally, once a rule is constructed, it's beneficial to re-use it because there's a bunch of allocations which definitely occur the first time that you compute a gradient, but which are typically not required on subsequent runs of the function unless you hit different control flow paths / change the size of arguments etc.

Do you have a particular use case in mind?