Skip to content

Type-stable rrule for applychain(::Vector) #1957

@ToucheSir

Description

@ToucheSir

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?

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions