JuliaDiff / JuliaDiff/ChainRulesCore.jl

Custom Covector Wrapper

Open
#165 1 comment 0 reactions 0 assignees View on GitHub

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^n and C^n -> R right first. Representing R->C^n as just complex vectors already works really well and I don't think we should change that. Zygote uses Adjoint for the C^n -> R, but this doesn't compose very well. The Problem here is that we want a real scalar product C^n x C^n = R^2n x R^2n -> R, so OneForm(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 a OneForm wrapper to wrap complex numbers a well as complex vectors, we could define such a composition.
  • n-dimensional arrays: We could also allow OneForm to wrap n-dimensional arrays to represent differentials of functions T = 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 product T x T -> R and 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 have OneForm(::Composite) as well. That way we also have a nice Composite x Composite -> R scalar 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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.