Description
Following up on some Slack (#autodiff) discussion with @mzgubic, I want to float the idea of adding a CRC rule "gallery" to the ChainRules docs, inspired by analogous things in the plotting world (eg BeautifulMakie).
I think that best practices for rule authors have become a bit too complicated to sufficiently demonstrate with just one or two small examples per section in the docs, particularly with interplay between
- custom structs & their constructors
- higher order functions
- CRC Thunks, Tangent types & AD callbacks
There are many clever rule authors in the wild who post helpful examples on Github/Discourse/Slack/Zulip but this information is very disaggregated and un-searchable. If there were some way for users to submit nice but less ultra-minimal examples and/or to aggregate relevant examples of custom rrule
and frule
definitions in published packages (as @ToucheSir suggested) I think linking to such a "gallery" from the CR docs would speed up the learning process for rule authors.
I would be glad to help realize such a thing if it seemed feasible and people wanted it.
Less organized thoughts from Slack:
A few custom rule examples I can find in forums but which I wouldn't know how to write based on CR docs:
- Constructors of widely used
AbstractArray
subtypes likeStaticArray
s,StructArray
s, etc. LinearAlgebra
methods of “implicit operator” types representing giant Matrices, eg from LinearMaps.jl- Iterative solver and partial decomposition methods involving the former, eg methods from IterativeSolvers.jl, KrylovKit.jl, Arpack.jl, etc, (“adjoint method” VJP rules)
- Higher order function methods defined outside Base, often versions of
map
orreduce
with parallelism. - Gradient accumulation into an
InplaceableThunk
of a custom struct