Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
ToucheSir opened this issue May 2, 2022 · 0 comments
Open

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

ToucheSir opened this issue May 2, 2022 · 0 comments

Comments

@ToucheSir
Copy link
Member

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?

@ToucheSir ToucheSir changed the title Type-stable rrule for applychain,(::Vector) Type-stable rrule for applychain(::Vector) May 2, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

No branches or pull requests

1 participant