JuliaDiff / JuliaDiff/ChainRules.jl
An additional approach to implementing rules
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 475
- Forks
- 98
- PR merge metrics
- No merged PRs in 30d
Description
There are basically two reasons to implement rules:
- to define AD. For example, you do have to tell an AD system somewhere how to differentiate addition and multiplication of floats,
- to make AD faster, without changing the semantics.
For 1 we obviously can't get around defining rules, however, for 2 we tend to implement rules in the same way as for 1 -- by completely over-riding any particular AD and just telling it how to differentiate a thing. However, one thing that we've not explored to a particularly great extent is re-writing code to make it more AD friendly, and then just saying "run AD on this".
Leaving aside concerns about the best way to achieve a code re-write for a minute, suppose that you wished to implement an rrule for *(::AbstractMatrix, ::Diagonal). LinearAlgebra implements this as follows:
(*)(A::AbstractMatrix, D::Diagonal) =
rmul!(copyto!(similar(A, promote_op(*, eltype(A), eltype(D.diag)), size(A)), A), D)
The problem from the perspective of a reverse-mode AD tool (that doesn't know how to handle mutation) is that the underlying implementation of this non-mutating operation is mutating. However, it is really quite clear how a non-mutating version of this operation could be implemented by looking at the definition of rmul!. Specifically, something like
A .* permutedims(D.diag)
This is the kind of code that we could plausible hope to run one of our current (or near-future) reverse-mode AD tools on, and have it do something sensible, whereas there was really no hope with LinearAlgebras definition.
Moreover, this kind of approach seems simpler for the rule-writer: rather than having to know how to differentiate a function, the rule-implementer just needs to know how to re-write the primal pass in a way that is more friendly towards AD.
This kind of approach is only valuable if there's functionality that could be easily re-written in an AD-friendly manner. My hypothesis is that there is lots of functionality in Base / the standard libraries that satisfies this because it was implemented by
- implementing a mutating version of a function (e.g.
gemm!) - implementing the non-mutating version of a function in terms of the mutating version. (e.g.
gemm)
This could provide a very simple partial solution to #232 by alleviating the need for generic rules in favour of code re-writes which are much more straightforward to achieve and lets the AD system auto-generate appropriate cotangents.
Thoughts on the principle? We can discuss implementation details once we've established whether or not we basically like the idea.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the linked Julia LinearAlgebra definitions in diagonal.jl, especially *(::AbstractMatrix, ::Diagonal) and rmul!, then read issue #232 for related context. This issue is currently a principle-level proposal; done would require maintainer agreement on whether rewrite-based rules should be pursued before implementation details are defined.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- tooling
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100