JuliaDiff / JuliaDiff/ChainRulesCore.jl
Ability to specify different rules based on what combinations of inputs are actually being used
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 267
- Forks
- 66
- PR merge metrics
- No merged PRs in 30d
Description
It would be nice to have an extension to the ChainRules API that allowed for different rules to be written and hit depending on what combination of inputs the derivative is being taken with respect to.
(or maybe similar analogy for forwards mode)
Basically partial derivative rules.
Thunking is a simple approximation to this (with it's own set of struggles)
@wsmoses requested for Enzyme, though this is also relevant to every kind of operator overloading based AD (since only tracked types etc will have derivatives taken wrt to them).
In contrast it is useless for Zygote/Diffractor as they do no kind of activity analysis etc, and transform absolutely all code that is run.
A bit of a sketch for what that API might look-like is in https://gist.github.com/oxinabox/c6ad25c468b3108f8a799bda66c147f8/
This might also be useful for partial mutation support, since it is probably completely safe to have rules for things that mutate inputs that are not "active" on the derivative path? (cf https://github.com/JuliaDiff/ChainRules.jl/pull/521)
Though as the main reason we don't do mutation is tied to Diffractor/Zygote not supporting it, that might be kinda moot, unless they got some at least some basic activity analysis.
(NB: we may not initially implement this in ChainRulesCore. It might be better to make a little experimental extension package for it first.)
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 by reviewing the ChainRules API and the linked gist, then compare the proposed extension with the existing ChainRulesCore design. Check the discussion around PR #521 and determine whether an experimental extension package should host the work; done means a decided API and implementation plan for input-combination-specific derivative rules.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend-api-design
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100