JuliaDiff / JuliaDiff/ChainRules.jl

Representing (co)tangents of structured matrices

Open
#191 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

design
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 Symmetric matrix should also produce a Symmetric matrix.
Is it also true that the pullback of a function that takes a Symmetric matrix should produce a Symmetric matrix?

@ChrisRackauckas:

yes

@sethaxen:

Okay, then what about when the input is Symmetric{<:Real}, but the pullback is passed an AbstractMatrix{<:Complex}? Should the pullback then produce a Hermitian{<:Complex} or a Symmetric{<: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 a Composite{<:Symmetric} ) to the pullback (e.g. a plain old AbstractMatrix ) 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 a Symmetric matrix as a general Matrix or 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

  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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.