JuliaDiff / JuliaDiff/ChainRules.jl
Representing (co)tangents of structured matrices
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 475
- Forks
- 98
- PR merge metrics
- No merged PRs in 30d
Description
This follows up a discussion on Slack.
@sethaxen:
The pushforward of a function that produces a
Symmetricmatrix should also produce aSymmetricmatrix.
Is it also true that the pullback of a function that takes aSymmetricmatrix should produce aSymmetricmatrix?
@ChrisRackauckas:
yes
@sethaxen:
Okay, then what about when the input is
Symmetric{<:Real}, but the pullback is passed anAbstractMatrix{<:Complex}? Should the pullback then produce aHermitian{<:Complex}or aSymmetric{<:Real}, or something else?
@willtebbutt:
I’ve been wondering about this.
My old answer would have been (assuming that the symmetric matrix looks at the upper triangular of the underlying matrix): add the lower triangle to the upper triangle (don’t touch the diagonal), and represent the cotangent via the upper triangle, the rationale being that the lower triangle of the data backing the symmetric matrix never gets used.
My newer thinking is that passing an asymmetric cotangent (I’m unclear whether we want represent this literally as a Symmetric or aComposite{<:Symmetric}) to the pullback (e.g. a plain oldAbstractMatrix) should just error. My reasoning for this is that if you’ve somehow managed to get an asymmetric cotangent, then the downstream operations have been implemented incorrectly (e.g.getindex(::Symmetric, ::Int, ::Int)must be wrong and treating aSymmetricmatrix as a generalMatrixor something). Nothing in the Julia ecosystem currently implements this AFAICT though.
There are several points for discussion here. Under this perspective (@willtebbutt's newer thinking, which I tend to agree with), if downstream rules and AD have done everything right, then the pullback for Y = Symmetric(A) should always receive an object ΔY with a data field (either Composite{<:Symmetric} or Symmetric), and its pullback should just be ΔY.data. If the pullback is passed an UpperTriangular, LowerTriangular or Diagonal matrix, as in the current rrule implementation and as in #178, then something is wrong somewhere else. Moreover, we don't need to do anything to data, such as zeroing a triangle, because if that triangle should be zeroed, it is already zeroed in ΔY.data (e.g. if Matrix(Y) was called, then the unused triangle was overwritten by the used triangle in the forward pass. Consequently, a correctly implemented Matrix_pullback will zero out the unused triangle in the cotangent vector before wrapping with Symmetric or Composite{Symmetric}). Thus a custom rule is probably not even necessary for the Symmetric constructor. Have I got this right?
One thing that worries me is e.g. what if a user defined an override like (::Diagonal * ::MyDiagonal)::MyDiagonal. This would trigger our general rrule. The pullback would expect an AbstractMatrix, but it will be passed a Composite{MyDiagonal} (or MyDiagonal). To do the right thing, it would need to produce a Composite{Diagonal} (or Diagonal) and a Composite{MyDiagonal} (or MyDiagonal). So how should we define generic rules that handle such cases?
Also, should we adopt a convention regarding whether the (co)tangent of structured matrices should be matrices or Composite? A point for the former is that we can automatically multiply them by other matrices and add them, and things should just work. A point for the latter is that in many cases, the (co)tangent doesn't share the same structure as the primal (e.g. the (co)tangent of a unitary matrix is a unitary transformation of a skew-Hermitian matrix). A compromise is a utility method that in most cases is a no-op but is meant to convert from a composite type to a primal when possible.
Relates https://github.com/JuliaDiff/ChainRules.jl/issues/52
cc @mcabbott @oxinabox
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 discussion in this issue, then review the related issues #178 and #52. The work requires deciding how structured matrix (co)tangents should be represented, how generic rules handle custom structured types, and whether a conversion convention is needed; it is done only after those design questions have a settled, tested direction.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- devtools
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 18/100