-
-
Notifications
You must be signed in to change notification settings - Fork 611
Description
Now that we have a concrete use case for it, I think this is worth considering. The rough structure would be similar to what Zygote automatically generates now. Namely, keep a stack of pullbacks (from rrule_via_ad
) around when running forward, and apply them in reverse in the pullback.
The main challenge lies in constraining the type of the stack and/or the final output. Push comes to shove I believe this is possible with return_type
, but ideally keeping track of a running overall pullback/output type would be sufficient. It may be simpler to only take this path if the eltype of layers
is concrete, but I'm not sure of how to do that with just dispatch.
Another consideration is double-differentiation, but IIUC we can use the opt-out functionality in CRC for this?