Compatibility with ChainRules, Drop Julia 1.0 support #414

rick2047 · 2021-04-14T17:00:20Z

related to #396. As mentioned by @moesphere in that thread, we need to define a rrule for the Interolations type. I've defined a simple rrule which enables me to integrate with, for example, Zygote like this:

julia> using Interpolations

julia> y = sin.(1:10)
10-element Vector{Float64}:
  0.8414709848078965
  0.9092974268256817
  0.1411200080598672
 -0.7568024953079282
 -0.9589242746631385
 -0.27941549819892586
  0.6569865987187891
  0.9893582466233818
  0.4121184852417566
 -0.5440211108893698

julia> itp = interpolate(y,BSpline(Cubic(Reflect(OnCell()))))
10-element interpolate(OffsetArray(::Vector{Float64}, 0:11), BSpline(Cubic(Reflect(OnCell())))) with element type Float64:
  0.7125302276762112
  0.769963450028415
  0.1194958272928704
 -0.6408357079724973
 -0.8119858486932346
 -0.23659994479000876
  0.556314857216602
  0.8377563450756788
  0.3489685127835064
 -0.29399432638595374

julia> Zygote.gradient(itp, 2)
([-0.35017548837401463],)

julia> Interpolations.gradient(itp, 2)
1-element StaticArrays.SVector{1, Float64} with indices SOneTo(1):
 -0.35017548837401463

This should allow us to integrate with Flux.

codecov · 2021-04-14T17:09:38Z

Codecov Report

Merging #414 (8a82a9d) into master (3b00087) will increase coverage by 0.12%.
The diff coverage is 100.00%.

❗ Current head 8a82a9d differs from pull request most recent head 1ca330c. Consider uploading reports for the commit 1ca330c to get more accurate results

@@            Coverage Diff             @@
##           master     #414      +/-   ##
==========================================
+ Coverage   84.40%   84.52%   +0.12%     
==========================================
  Files          23       24       +1     
  Lines        1590     1596       +6     
==========================================
+ Hits         1342     1349       +7     
+ Misses        248      247       -1

Impacted Files	Coverage Δ
src/Interpolations.jl	`76.98% <ø> (+0.18%)`	⬆️
src/chainrules/chainrules.jl	`100.00% <100.00%> (ø)`
src/b-splines/indexing.jl	`64.22% <0.00%> (+0.81%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9375e3b...1ca330c. Read the comment docs.

mkitti · 2021-04-14T21:28:22Z

Awesome. What's the best way to test this?

rick2047 · 2021-04-15T06:12:20Z

Actually I was going to ask the community the same question. I can of course write a naive test which takes gradient over a known function and we compare values. But that isn't the best approach I think. What we really want to do is check integration with chain rules.

The discussion started with trying to implement upsample methods. Maybe we can do something with that?

rick2047 · 2021-04-15T06:18:21Z

I found ChainRulesTestUtils.jl I guess that's the best approach.

rick2047 · 2021-04-15T21:24:27Z

I tried using ChainRulesTestUtils but it seems like it doesn't work for Interpolation objects. Adding code here in case I am doing something obviously wrong.

julia> using ChainRulesTestUtils

julia> x = collect(1:10)
julia> y = sin.(x)

julia> itp = interpolate(y,BSpline(Linear()))


julia> test_rrule(itp, 1, 1)
test_rrule: [0.8414709848078965, 0.9092974268256817, 0.1411200080598672, -0.7568024953079282, -0.9589242746631385, -0.27941549819892586, 0.6569865987187891, 0.9893582466233818, 0.4121184852417566, -0.5440211108893698] at (1, 1): Error During Test at C:\Users\paresh\.julia\packages\ChainRulesTestUtils\NCO8i\src\testers.jl:168
  Got exception outside of a @test
  ArgumentError: test_rrule cannot be used on closures/functors (such as [0.8414709848078965, 0.9092974268256817, 0.1411200080598672, -0.7568024953079282, -0.9589242746631385, -0.27941549819892586, 0.6569865987187891, 0.9893582466233818, 0.4121184852417566, -0.5440211108893698])
  Stacktrace:
    [1] _ensure_not_running_on_functor(f::Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Linear}, Tuple{Base.OneTo{Int64}}}, name::String)
      @ ChainRulesTestUtils ~\.julia\packages\ChainRulesTestUtils\NCO8i\src\testers.jl:238
    [2] macro expansion
      @ ~\.julia\packages\ChainRulesTestUtils\NCO8i\src\testers.jl:169 [inlined]
    [3] macro expansion
      @ C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.6\Test\src\Test.jl:1151 [inlined]
    [4] test_rrule(::Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Linear}, Tuple{Base.OneTo{Int64}}}, ::Int64, ::Vararg{Int64, N} 
where N; output_tangent::ChainRulesTestUtils.Auto, fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, check_inferred::Bool, fkwargs::NamedTuple{(), Tuple{}}, rtol::Float64, atol::Float64, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
      @ ChainRulesTestUtils ~\.julia\packages\ChainRulesTestUtils\NCO8i\src\testers.jl:169
    [5] test_rrule(::Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Linear}, Tuple{Base.OneTo{Int64}}}, ::Int64, ::Vararg{Int64, N} 
where N)
      @ ChainRulesTestUtils ~\.julia\packages\ChainRulesTestUtils\NCO8i\src\testers.jl:166
    [6] top-level scope
      @ REPL[32]:1
    [7] eval
      @ .\boot.jl:360 [inlined]
    [8] eval
      @ .\Base.jl:39 [inlined]
    [9] repleval(m::Module, code::Expr, #unused#::String)
      @ VSCodeServer ~\.vscode\extensions\julialang.language-julia-1.1.38\scripts\packages\VSCodeServer\src\repl.jl:124
   [10] (::VSCodeServer.var"#47#49"{Module, Expr, REPL.LineEditREPL, REPL.LineEdit.Prompt})()
      @ VSCodeServer ~\.vscode\extensions\julialang.language-julia-1.1.38\scripts\packages\VSCodeServer\src\repl.jl:99
   [11] with_logstate(f::Function, logstate::Any)
      @ Base.CoreLogging .\logging.jl:491
   [12] with_logger
      @ .\logging.jl:603 [inlined]
   [13] (::VSCodeServer.var"#46#48"{Module, Expr, REPL.LineEditREPL, REPL.LineEdit.Prompt})()
      @ VSCodeServer ~\.vscode\extensions\julialang.language-julia-1.1.38\scripts\packages\VSCodeServer\src\repl.jl:100
   [14] #invokelatest#2
      @ .\essentials.jl:708 [inlined]
   [15] invokelatest(::Any)
      @ Base .\essentials.jl:706
   [16] macro expansion
      @ ~\.vscode\extensions\julialang.language-julia-1.1.38\scripts\packages\VSCodeServer\src\eval.jl:34 [inlined]
   [17] (::VSCodeServer.var"#60#61")()
      @ VSCodeServer .\task.jl:406
Test Summary:
                                                                      | Error  Total
test_rrule: [0.8414709848078965, 0.9092974268256817, 0.1411200080598672, -0.7568024953079282, -0.9589242746631385, -0.27941549819892586, 0.6569865987187891, 0.9893582466233818, 0.4121184852417566, -0.5440211108893698] at (1, 1) |     1      1
ERROR: Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.

So this seems like a dud, any other ideas about testing are welcome

DhairyaLGandhi · 2021-04-16T08:16:18Z

src/chainrules/chainrules.jl

+function ChainRulesCore.rrule(itp::AbstractInterpolation, x...)
+    y = itp(x...)
+    function back(Δ)
+        (NO_FIELDS, Δ * gradient(itp, x...))


Would you not want to update any of the fields of the interpolation too?

I am not sure about that. I am confident that the second argument needs to be the gradient, but I am not sure what the first argument is even supposed to be. I have hacked this in to see if it works but I am not confident this is right. What would you recommend? Maybe you can clarify a bit more on what the first and second argument are supposed to be

I will change the call to Interpolations.gradient, I made the choice at random anyways.

DhairyaLGandhi · 2021-04-16T08:17:57Z

Yay! This looks like a really nice step. Maybe qualify the call to gradient as Interpolations.gradient for readability if nothing else?

DhairyaLGandhi · 2021-04-16T08:18:49Z

I would directly test that zygote is producing the correct and expected value

rick2047 · 2021-04-16T14:31:51Z

@DhairyaLGandhi thanks for your attention here. I have added a naive test, if this works we can think about adding more comprehensive testing.

rick2047 · 2021-04-16T14:36:14Z

I dug around a bit more on ChainRulesTestUtils, they already have an issue open for the functor issue JuliaDiff/ChainRulesTestUtils.jl#117 (comment) . It seems like this is already a milestone for their v1. I would still proceed with the zygote but eventually we should shift to this package

rick2047 · 2021-04-16T15:19:25Z

I have no idea why these tests are failing. Did I do dependencies wrong?

mkitti · 2021-04-16T15:26:08Z

Interpolations: Error During Test at D:\a\Interpolations.jl\Interpolations.jl\test\chainrules.jl:10
  Test threw exception
  Expression: (Zygote.gradient(itp, 1))[1] == Interpolations.gradient(itp, 1)
  ArgumentError: unable to check bounds for indices of type Interpolations.WeightedAdjIndex{2,Int64}

rick2047 · 2021-04-16T15:30:22Z

@mkitti I saw that, that is the original error we started to solve. I don't understand how it works on 3 tests and doesn't work on 3 tests.

also, itworksonmymachine.jpg :(

mkitti · 2021-04-16T15:38:19Z

Ah I see. The tests are failing on Julia 1.0.5. They succeed on Julia 1.6.

mkitti · 2021-04-16T16:01:53Z

Where the tests are failing:

 Installed Zygote ─────────────── v0.4.20
 Installed ZygoteRules ────────── v0.2.1

…ns.jl into AD_interop

rick2047 · 2021-04-17T09:36:32Z

I added another test for 2D interpolation case. Now I am a bit stuck about how should we proceed. @DhairyaLGandhi is still scrutinizing the actual rule, but apart from that it seems like integration with Zygote is working. Ideas on how to proceed are welcome, I was thinking maybe we can write what is asked for in this issue JuliaImages/ImageTransformations.jl#113?

DhairyaLGandhi · 2021-04-17T10:26:57Z

What else is needed right now? I am thinking it might be good to merge this and move from there. What do you think @mkitti ?

mkitti · 2021-04-17T10:47:52Z

The new code is in a mergeable state. Some housekeeping is required.

If we're going to break compat with Julia 1.0, this probably should be v0.14.0. However, it's only clear to me that we need Julia 1.3 for testing.

Do we have any bounds on compat with ChainRulesCore.jl or should we just declare compat against the latest version?

DhairyaLGandhi · 2021-04-17T11:26:01Z

The current minor version of ChainRulesCore should be fine. You can also make a patch release, it's allowed within semver.

rick2047 · 2021-04-17T15:00:24Z

If this is mergeable thats do it. this doesn't break anything and adds new functionality. The only apprehension I had was if our rule, which basically is Zygote.gradient(itp,x)[1] == Interpolations.gradient(itp,x) the right definition. Should I document this somewhere?

mkitti · 2021-04-17T15:10:30Z

Well the tests do break Julia 1.0 compatibility, which is fine. It just needs to be managed. I'll probably run a v0.13 branch that is compatible with Julia 1.0 for a while longer in parallel with a v0.14 branch that requires Julia 1.3.

Documentation would be great. Could you add a docs/src/chainrules.md to describe this?

…ns.jl into AD_interop

DhairyaLGandhi · 2021-04-17T16:26:40Z

src/chainrules/chainrules.jl

+function ChainRulesCore.rrule(itp::AbstractInterpolation, x...)
+    y = itp(x...)
+    function pullback(Δy)
+        (NO_FIELDS, Δy * Interpolations.gradient(itp, x...))


One quick thought (sorry this should have come earlier) NO_FIELDS -> DoesNotExist(). @rick2047 would you be up for a tiny follow up?

sure, I was bothered by this too. but this is already merged, can I still push to this branch?

You can use the same branch. It's your branch. Just send a new PR when you are ready.

mkitti · 2021-06-08T05:02:37Z

We seem to be having CI issues with ChainRulesCore.jl 0.10. @rick2047 @DhairyaLGandhi , do you have any ideas where the issue may be?

naive implementation of rrule for chainrules compatibility

d33b9e8

DhairyaLGandhi reviewed Apr 16, 2021

View reviewed changes

rick2047 added 2 commits April 16, 2021 16:20

change function call to make it more explicite

e21fd68

POC of testing

b236c9c

mkitti and others added 3 commits April 16, 2021 12:06

Test on Julia 1.3 and 1

f3f001e

make testset and add a 2D example

030a501

Merge branch 'AD_interop' of https://github.com/rick2047/Interpolatio…

ef2ff6c

…ns.jl into AD_interop

Set ChainRulesCore compat to 0.9, Julia compat 1.3

e085a0f

rick2047 added 2 commits April 17, 2021 17:47

add docstring, chaing pullback function name

3485887

add separate doc explaining some rationales

4debd36

rick2047 and others added 2 commits April 17, 2021 17:47

Merge branch 'AD_interop' of https://github.com/rick2047/Interpolatio…

8a82a9d

…ns.jl into AD_interop

Merge branch 'master' into AD_interop

1ca330c

mkitti changed the title ~~[WIP] Compatibility with ChainRules~~ Compatibility with ChainRules Apr 17, 2021

mkitti changed the title ~~Compatibility with ChainRules~~ Compatibility with ChainRules, Drop Julia 1.0 support Apr 17, 2021

mkitti merged commit ee4a49b into JuliaMath:master Apr 17, 2021

DhairyaLGandhi reviewed Apr 17, 2021

View reviewed changes

This was referenced Apr 17, 2021

Bug fix in ChainRulesCore pullback function #416

Merged

compatibility with Flux #396

Open

johnnychen94 mentioned this pull request Jun 28, 2021

Interpolations v0.14? #433

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Compatibility with ChainRules, Drop Julia 1.0 support #414

Compatibility with ChainRules, Drop Julia 1.0 support #414

rick2047 commented Apr 14, 2021 •

edited

Loading

codecov bot commented Apr 14, 2021 •

edited

Loading

mkitti commented Apr 14, 2021

rick2047 commented Apr 15, 2021

rick2047 commented Apr 15, 2021

rick2047 commented Apr 15, 2021

DhairyaLGandhi Apr 16, 2021

rick2047 Apr 16, 2021

rick2047 Apr 16, 2021

DhairyaLGandhi commented Apr 16, 2021

DhairyaLGandhi commented Apr 16, 2021

rick2047 commented Apr 16, 2021

rick2047 commented Apr 16, 2021 •

edited

Loading

rick2047 commented Apr 16, 2021

mkitti commented Apr 16, 2021

rick2047 commented Apr 16, 2021

mkitti commented Apr 16, 2021

mkitti commented Apr 16, 2021

rick2047 commented Apr 17, 2021

DhairyaLGandhi commented Apr 17, 2021

mkitti commented Apr 17, 2021

DhairyaLGandhi commented Apr 17, 2021

rick2047 commented Apr 17, 2021

mkitti commented Apr 17, 2021

DhairyaLGandhi Apr 17, 2021

rick2047 Apr 17, 2021

mkitti Apr 17, 2021

mkitti commented Jun 8, 2021

Compatibility with ChainRules, Drop Julia 1.0 support #414

Compatibility with ChainRules, Drop Julia 1.0 support #414

Conversation

rick2047 commented Apr 14, 2021 • edited Loading

codecov bot commented Apr 14, 2021 • edited Loading

Codecov Report

mkitti commented Apr 14, 2021

rick2047 commented Apr 15, 2021

rick2047 commented Apr 15, 2021

rick2047 commented Apr 15, 2021

DhairyaLGandhi Apr 16, 2021

Choose a reason for hiding this comment

rick2047 Apr 16, 2021

Choose a reason for hiding this comment

rick2047 Apr 16, 2021

Choose a reason for hiding this comment

DhairyaLGandhi commented Apr 16, 2021

DhairyaLGandhi commented Apr 16, 2021

rick2047 commented Apr 16, 2021

rick2047 commented Apr 16, 2021 • edited Loading

rick2047 commented Apr 16, 2021

mkitti commented Apr 16, 2021

rick2047 commented Apr 16, 2021

mkitti commented Apr 16, 2021

mkitti commented Apr 16, 2021

rick2047 commented Apr 17, 2021

DhairyaLGandhi commented Apr 17, 2021

mkitti commented Apr 17, 2021

DhairyaLGandhi commented Apr 17, 2021

rick2047 commented Apr 17, 2021

mkitti commented Apr 17, 2021

DhairyaLGandhi Apr 17, 2021

Choose a reason for hiding this comment

rick2047 Apr 17, 2021

Choose a reason for hiding this comment

mkitti Apr 17, 2021

Choose a reason for hiding this comment

mkitti commented Jun 8, 2021

rick2047 commented Apr 14, 2021 •

edited

Loading

codecov bot commented Apr 14, 2021 •

edited

Loading

rick2047 commented Apr 16, 2021 •

edited

Loading