chain rules for DCT #273

vpuri3 · 2023-06-26T01:51:19Z

address #272

codecov · 2023-06-26T01:54:24Z

Codecov Report

Attention: 18 lines in your changes are missing coverage. Please review.

Comparison is base (ef8fc5b) 73.08% compared to head (26df888) 70.70%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #273      +/-   ##
==========================================
- Coverage   73.08%   70.70%   -2.38%     
==========================================
  Files           5        6       +1     
  Lines         535      553      +18     
==========================================
  Hits          391      391              
- Misses        144      162      +18

Files	Coverage Δ
src/FFTW.jl	`85.71% <ø> (ø)`
ext/FFTWChainRulesCoreExt.jl	`0.00% <0.00%> (ø)`

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

vpuri3 · 2023-06-26T21:57:10Z

This error is only happening with the MKL provider. With MKL, FFTW.jl doesn't even compile on my machine. Could be due to 008bc5b?

test_frule: idct on Array{Float64, 3},Int64: Error During Test at /home/runner/.julia/packages/ChainRulesTestUtils/C9L2i/src/testers.jl:123
  Got exception outside of a @test
  FFTW could not create plan
  Stacktrace:
    [1] error(s::String)
      @ Base ./error.jl:35
    [2] macro expansion
      @ FFTW ~/work/FFTW.jl/FFTW.jl/src/fft.jl:722 [inlined]

vpuri3 · 2023-07-08T00:40:58Z

@devmotion, could you please review this? LMK if you want me to remove the requires part.

devmotion

Can you explain the idea of the PR? The design goal was to define the differentiation rules in AbstractFFTs via adjoint plans. Downstream packages were supposed to implement this new adjoint interface.

vpuri3 · 2023-07-08T00:44:35Z

Only the MKL tests are failing. Is there a regression with MKL?

devmotion · 2023-07-08T00:47:30Z

Can you explain the idea of the PR? The design goal was to define the differentiation rules in AbstractFFTs via adjoint plans. Downstream packages were supposed to implement this new adjoint interface.

More concretely, #249 outlines the intended approach. FFTW was not supposed to define custom rules.

vpuri3 · 2023-07-08T00:49:24Z

This PR defines Chain Rules for FFTW.dct, idct since these methods are not defined in AbstractFFTs.

devmotion · 2023-07-08T00:55:04Z

Ref: JuliaMath/AbstractFFTs.jl#56

vpuri3 · 2023-07-08T00:55:15Z

With #249 , gradient computation would error for DCT/IDCT. This is because dct(x) is defined as dct(x) = plan_dct(x) * x, and Zygote ends up going inside plan_dct where it encounters some error. This is why we need an rrule for dct.

julia> using FFTW, Zygote                                                                                                                                                                                       
                                                                                                                                                                                                                
julia>                                                                                                                                                                                                          
julia> using LinearAlgebra, FFTW, Zygote                                                                                                                                                                        
                                                                                                                                                                                                                
julia> x = rand(4)                                                                                                                                                                                              
4-element Vector{Float64}:                                                                                                                                                                                      
 0.8692266334693106                                                                                                                                                                                             
 0.6938635624794242                                                                                                                                                                                             
 0.552208368655668                                                                                                                                                                                              
 0.9197557963740512                                                                                                                                                                                             
                                                                                                                                                                                                                
julia> f(x) = x |> dct |> idct |> norm                                                                                                                                                                          
f (generic function with 1 method)                                                                                                                                                                              
                                                                                                                                                                                                                
julia> f(x)                                                                                                                                                                                                     
1.5452787421840921                                                                                                                                                                                              
                                                                                                                                                                                                                
julia> Zygote.gradient(f, x)                                                                                                                                                                                    
ERROR: Compiling Tuple{Type{FFTW.r2rFFTWPlan{Float64, Any, false, 1}}, Vector{Float64}, FFTW.FakeArray{Float64, 1}, UnitRange{Int64}, Int64, UInt32, Float64}: try/catch is not supported.                      
Refer to the Zygote documentation for fixes.        
https://fluxml.ai/Zygote.jl/latest/limitations

devmotion

Sorry, somehow I missed yesterday that the PR does not add rules for plans but for the dct, idct and r2r functions 🤦

As long as they/their interface is not moved to AbstractFFTs, rules should be defined here 👍

Project.toml

ext/FFTWChainRulesCoreExt.jl

devmotion · 2023-07-08T18:34:00Z

ext/FFTWChainRulesCoreExt.jl

+
+# R2R
+
+function ChainRulesCore.frule(Δ, ::typeof(r2r), x::AbstractArray, region...)


It seems the rrule for r2r is missing?

The R2R transforms are not unitary. There is some scaling involved that depends on the kind of R2R transform. Because it looks like an involved task, I chose to skip that for now. I am happy to look into that in a separate PR

src/FFTW.jl

test/runtests.jl

Co-authored-by: David Widmann <[email protected]>

vpuri3 · 2023-07-10T20:41:02Z

@devmotion I've addressed all your comments. LMK if you have more questions :D

vpuri3 · 2023-07-15T13:09:14Z

@devmotion ping :)

vpuri3 added 2 commits June 25, 2023 21:49

chainrulescore ext

7abaeeb

add dct chainrules tests

1a9cd68

add CRC to test dependencies

d87c550

vpuri3 added 4 commits June 26, 2023 19:12

rules for r2r WIP

1f321f8

rename args -> region

8da2495

rename dims -> region

5a975ab

rm r2r rrule. needs more thinking

39c7b6d

devmotion reviewed Jul 8, 2023

View reviewed changes

vpuri3 and others added 8 commits July 8, 2023 16:03

Update Project.toml

10ac31c

Co-authored-by: David Widmann <[email protected]>

Update src/FFTW.jl

3ee37e7

Co-authored-by: David Widmann <[email protected]>

rm requires

74b8c75

uncomment tests

a31b52d

mv include(CRCext) to bottom to fix precomile failure in <1.9

3d6e160

default region to 1:ndims(x) in frules

d0ef838

different rrules for dct(x), dct(x, region)

6b8d178

put DCT/IDCT rules in a loop

608aa02

vpuri3 mentioned this pull request Jul 9, 2023

Tracking issue for plan chain rule support JuliaMath/AbstractFFTs.jl#94

Open

5 tasks

Merge branch 'master' into dct

26df888

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chain rules for DCT #273

chain rules for DCT #273

vpuri3 commented Jun 26, 2023 •

edited

Loading

codecov bot commented Jun 26, 2023 •

edited

Loading

vpuri3 commented Jun 26, 2023

vpuri3 commented Jul 8, 2023

devmotion left a comment

vpuri3 commented Jul 8, 2023

devmotion commented Jul 8, 2023

vpuri3 commented Jul 8, 2023

devmotion commented Jul 8, 2023

vpuri3 commented Jul 8, 2023 •

edited

Loading

devmotion left a comment

devmotion Jul 8, 2023

vpuri3 Jul 8, 2023

vpuri3 commented Jul 10, 2023 •

edited

Loading

vpuri3 commented Jul 15, 2023


		# R2R

		function ChainRulesCore.frule(Δ, ::typeof(r2r), x::AbstractArray, region...)

chain rules for DCT #273

Are you sure you want to change the base?

chain rules for DCT #273

Conversation

vpuri3 commented Jun 26, 2023 • edited Loading

codecov bot commented Jun 26, 2023 • edited Loading

Codecov Report

vpuri3 commented Jun 26, 2023

vpuri3 commented Jul 8, 2023

devmotion left a comment

Choose a reason for hiding this comment

vpuri3 commented Jul 8, 2023

devmotion commented Jul 8, 2023

vpuri3 commented Jul 8, 2023

devmotion commented Jul 8, 2023

vpuri3 commented Jul 8, 2023 • edited Loading

devmotion left a comment

Choose a reason for hiding this comment

devmotion Jul 8, 2023

Choose a reason for hiding this comment

vpuri3 Jul 8, 2023

Choose a reason for hiding this comment

vpuri3 commented Jul 10, 2023 • edited Loading

vpuri3 commented Jul 15, 2023

vpuri3 commented Jun 26, 2023 •

edited

Loading

codecov bot commented Jun 26, 2023 •

edited

Loading

vpuri3 commented Jul 8, 2023 •

edited

Loading

vpuri3 commented Jul 10, 2023 •

edited

Loading