Add adjoints for functions on Hermitian matrices (take 2) #355

sethaxen · 2019-10-03T00:09:20Z

Note: this is a redo of #352 because of a wayward rebase

This PR adds adjoints for the following functions of Hermitian and real symmetric matrices:

Each of these can be written as a matrix power series, which for a diagonalizable matrix simplifies to element-wise action of the function upon the eigenvalues of the matrix. Consequently, their computation in LinearAlgebra and their adjoints take a similar form to each other.

Along the way, I fixed #340 for generic matrices, since the current adjoint for exp using the eigendecomposition shares a similar form to the Hermitian cases.

sethaxen · 2019-10-07T20:46:56Z

The test is currently failing because DiffRule's complex derivative of acosh is wrong. Looks like this PR fixes it but hasn't been worked on in a while. Should I just add a custom acosh adjoint to Zygote?

MikeInnes · 2019-10-08T12:26:10Z

Should I just add a custom acosh adjoint to Zygote?

Sure, please do.

This reverts commit dbbf12f.

This reverts commit 54c1c4b.

1.0 doesn't have LinearAlgebra.RealHermSym

All of the functions considered here have the same domain as range.

Post-processing just enforces type constraints, but pulling back through it significantly changes gradients wrt args. It turns out this isn't necessary.

Enables overloading for specific functions

Separating real and imag makes gradients discontinuous when output type is determined by input value.

sethaxen · 2019-10-10T07:08:29Z

The ^(A, p::Integer) adjoints are still erroring due to #344, but otherwise, this is ready for review.

src/lib/array.jl

When 2nd derivative is not available, use finite difference approximation of 2nd deriv for greater accuracy. When it is, relax the check.

sethaxen · 2019-10-23T19:45:16Z

Just a reminder that the only reason this doesn't pass is because Array(::Union{Symmetric,Hermitian}) mutates. I'd like to push this PR over the finish line, which I can do by either adding an adjoint for this case or if a solution like those discussed in (#344 and #377) is available.

willtebbutt

I'm not amazingly well-equipped to review this / am currently overloaded to don't have time properly dive into this. That being said, from what I can tell fairly quickly it LGTM -- everything is tested as far as I can tell. As such, I don't want to hold this up, and I can't see any reason not to merge.

MikeInnes · 2019-11-04T14:10:45Z

Since this is still marked WIP I'll give @sethaxen the ability to merge, rather than merging immediately. I'm not sure if any other help is needed from our side but lmk if so.

bors d+

bors · 2019-11-04T14:10:46Z

✌️ sethaxen can now approve this pull request. To approve and merge a pull request, simply reply with bors r+. More detailed instructions are available here.

sethaxen · 2019-11-04T15:52:38Z

bors try

bors · 2019-11-04T15:52:41Z

try

Merge conflict

sethaxen · 2019-11-04T22:57:21Z

@MikeInnes what's missing is this:

Just a reminder that the only reason this doesn't pass is because Array(::Union{Symmetric,Hermitian}) mutates. I'd like to push this PR over the finish line, which I can do by either adding an adjoint for this case or if a solution like those discussed in (#344 and #377) is available.

I'd appreciate feedback. I'd rather not merge something that fails the tests.

MikeInnes · 2019-11-08T12:31:38Z

Adding an adjoint for that case seems fine. We can remove it as part of #344 as needed, but that doesn't need to block this work.

sethaxen · 2019-11-09T23:08:07Z

bors try

bors · 2019-11-09T23:17:35Z

try

Build succeeded

ci/gitlab/gitlab.com

sethaxen · 2019-11-09T23:22:56Z

bors r+

355: [WIP] Add adjoints for functions on Hermitian matrices (take 2) r=sethaxen a=sethaxen _Note: this is a redo of #352 because of a wayward rebase_ This PR adds adjoints for the following functions of Hermitian and real symmetric matrices: - [x] `exp` - [x] `cos` - [x] `sin` - [x] `tan` - [x] `cosh` - [x] `sinh` - [x] `tanh` - [x] `log` - [x] `acos` - [x] `asin` - [x] `atan` - [x] `acosh` - [x] `asinh` - [x] `atanh` - [x] `sqrt` - [x] `sincos` - [x] `^` Each of these can be written as a matrix power series, which for a diagonalizable matrix simplifies to element-wise action of the function upon the eigenvalues of the matrix. Consequently, their computation in `LinearAlgebra` and their adjoints take a similar form to each other. Along the way, I fixed #340 for generic matrices, since the current adjoint for `exp` using the eigendecomposition shares a similar form to the Hermitian cases. Co-authored-by: Seth Axen <[email protected]>

bors · 2019-11-09T23:32:49Z

Build succeeded

ci/gitlab/gitlab.com

sethaxen added 9 commits October 8, 2019 09:01

Add adjoints real functions of hermitian matrices

cc5b0df

Make exp work for defective matrices. Fixes FluxML#340

ee1255e

Support arbitrary functions

69eef19

Unify Symmetric/Hermitian tests

e47eecb

Simplify hermitian eigen/eigvals tests

fe6ed06

Revert "Add adjoints for real, conj, and imag for arrays"

54c1c4b

This reverts commit dbbf12f.

Fix complex derivative for acosh

a9acf7e

Add sincos

f4e8871

Check if array is square

29db72f

sethaxen force-pushed the sdaxen/hermitian_funcs branch from 97b4471 to 29db72f Compare October 8, 2019 16:30

sethaxen added 7 commits October 8, 2019 09:35

Revert "Revert "Add adjoints for real, conj, and imag for arrays""

a513195

This reverts commit 54c1c4b.

Fix test for Julia 1.0

9e9a370

1.0 doesn't have LinearAlgebra.RealHermSym

Clean up notation

75afe81

Add adjoint for ^

bfb2d4d

Add type annotation

16d7271

All of the functions considered here have the same domain as range.

Avoid eigendecomposition for integer power

b1b88f9

Clean up notation

a1b10ec

sethaxen mentioned this pull request Oct 9, 2019

Add adjoints for Matrix and Vector #344

Closed

sethaxen added 10 commits October 9, 2019 16:14

Rename functions

5241fba

Add custom adjoint for _realifydiag!

b13d6d8

Don't pullback through post-processing

2841ed0

Post-processing just enforces type constraints, but pulling back through it significantly changes gradients wrt args. It turns out this isn't necessary.

Factor out computation of scalar func pullback

74225cf

Enables overloading for specific functions

Ensure output types are correct

1deb189

Don't split real and imag when output type variable

e377fae

Separating real and imag makes gradients discontinuous when output type is determined by input value.

Expand and stabilize tests

df8d3e6

Combine functions

c6942ef

Cover all cases in tests

378955e

Add custom pullback for exp

31cab4e

antoine-levitt reviewed Oct 10, 2019

View reviewed changes

src/lib/array.jl Outdated Show resolved Hide resolved

sethaxen added 2 commits October 10, 2019 23:37

Fix test for p=0

53b7f8e

Improve check and accuracy

62b5b22

When 2nd derivative is not available, use finite difference approximation of 2nd deriv for greater accuracy. When it is, relax the check.

MikeInnes force-pushed the master branch from 0d5df1b to 8f30e9d Compare October 14, 2019 12:16

MikeInnes requested a review from willtebbutt October 24, 2019 11:10

willtebbutt approved these changes Oct 27, 2019

View reviewed changes

Merge branch 'master' into sdaxen/hermitian_funcs

d8decf2

Add adjoints for converting Hermitian/Symmetric to matrix

fb95ec7

bors bot added a commit that referenced this pull request Nov 9, 2019

Try #355:

287f704

bors bot merged commit fb95ec7 into FluxML:master Nov 9, 2019

sethaxen deleted the sdaxen/hermitian_funcs branch November 9, 2019 23:33

sethaxen changed the title ~~[WIP] Add adjoints for functions on Hermitian matrices (take 2)~~ Add adjoints for functions on Hermitian matrices (take 2) Nov 12, 2019

sethaxen mentioned this pull request Nov 12, 2019

Handle degenerate eigenvalues in JVP rule for eigh jax-ml/jax#1665

Draft

sethaxen mentioned this pull request Dec 2, 2019

WIP: add adjoint for eigen #327

Closed

sethaxen mentioned this pull request May 17, 2020

Symmetric/Hermitian matrix function rules JuliaDiff/ChainRules.jl#193

Merged

1 task

shoyer mentioned this pull request Jun 16, 2020

support jax.numpy.linalg.eig eigenvector derivatives jax-ml/jax#2748

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add adjoints for functions on Hermitian matrices (take 2) #355

Add adjoints for functions on Hermitian matrices (take 2) #355

sethaxen commented Oct 3, 2019 •

edited

Loading

sethaxen commented Oct 7, 2019 •

edited

Loading

MikeInnes commented Oct 8, 2019

sethaxen commented Oct 10, 2019

sethaxen commented Oct 23, 2019

willtebbutt left a comment

MikeInnes commented Nov 4, 2019

bors bot commented Nov 4, 2019

sethaxen commented Nov 4, 2019

bors bot commented Nov 4, 2019

sethaxen commented Nov 4, 2019 •

edited

Loading

MikeInnes commented Nov 8, 2019

sethaxen commented Nov 9, 2019

bors bot commented Nov 9, 2019

sethaxen commented Nov 9, 2019

bors bot commented Nov 9, 2019

Add adjoints for functions on Hermitian matrices (take 2) #355

Add adjoints for functions on Hermitian matrices (take 2) #355

Conversation

sethaxen commented Oct 3, 2019 • edited Loading

sethaxen commented Oct 7, 2019 • edited Loading

MikeInnes commented Oct 8, 2019

sethaxen commented Oct 10, 2019

sethaxen commented Oct 23, 2019

willtebbutt left a comment

Choose a reason for hiding this comment

MikeInnes commented Nov 4, 2019

bors bot commented Nov 4, 2019

sethaxen commented Nov 4, 2019

bors bot commented Nov 4, 2019

try

Merge conflict

sethaxen commented Nov 4, 2019 • edited Loading

MikeInnes commented Nov 8, 2019

sethaxen commented Nov 9, 2019

bors bot commented Nov 9, 2019

try

Build succeeded

sethaxen commented Nov 9, 2019

bors bot commented Nov 9, 2019

Build succeeded

sethaxen commented Oct 3, 2019 •

edited

Loading

sethaxen commented Oct 7, 2019 •

edited

Loading

sethaxen commented Nov 4, 2019 •

edited

Loading