JuliaDiff / JuliaDiff/ChainRulesCore.jl
Custom Covector Wrapper
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 267
- Forks
- 66
- PR merge metrics
- No merged PRs in 30d
Description
What is becoming more and more apparent to me in light of #159 and #160 that neither Base.adjoint, nor Base.transpose are the right abstraction for covector differentials, so I have been tinkering a bit with the idea of rolling our own wrapper instead. To keep with adopting differential geometry terminology, we could name such a type OneForm, but I am open to a name that's more approachable to non-math folk. This could represent multiple different cases of directional derivatives.
- Complex Numbers: We discovered that Wirtinger derivatives are quite difficult to deal with, so I think we should focus on getting
R -> C^nandC^n -> Rright first. RepresentingR->C^nas just complex vectors already works really well and I don't think we should change that. Zygote usesAdjointfor theC^n -> R, but this doesn't compose very well. The Problem here is that we want a real scalar productC^n x C^n = R^2n x R^2n -> R, soOneForm(v)takes in a seed as a complex vector and spits out a real number, i.e. the directional derivative wrt the seed. The conjugation that Zygote uses gives us a complex scalar product, which is not really that useful to us. If we allowed aOneFormwrapper to wrap complex numbers a well as complex vectors, we could define such a composition. - n-dimensional arrays: We could also allow
OneFormto wrap n-dimensional arrays to represent differentials of functionsT = R^(m1 x m2 x ... x mn) ->R. AFAIK Zygote currently isn't very consistent about handling these, I believe for matrices it usually takes the conjugate transpose, but not for higher dimensional arrays. This way we could also define a scalar productT x T -> Rand it would compose well with forward diff. Composite: Composite are also basically just real vectors with some additional information, so I think it would also make sense to haveOneForm(::Composite)as well. That way we also have a niceComposite x Composite -> Rscalar product relation
I sketched out a prototype in the sim/one_form branch, but I am still not quite happy with it. Uni currently keeps me fairly busy, but hope I can dedicate some time after next month. My worry is that ChainRules is already quite big and adopting something like this would have to touch pretty much everything, but I think it could be worth it. Do people generally agree with this solution, or should we go in a different direction? I know @willtebbutt is also interested in this, this doesn't yet solve the C -> C (holomorphic) case, but I believe it would let us handle the most common cases, I would love to hear other thoughts on this!
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 issues #159 and #160, then inspect the prototype in the sim/one_form branch. Compare the proposed wrapper across complex numbers, n-dimensional arrays, and Composite values, and seek agreement on the abstraction and scope before implementation. Done is not defined yet because the issue is an open design discussion.
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
- 20/100