Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix periodic kernel AD #531

Merged
merged 21 commits into from
Feb 7, 2024
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/KernelFunctions.jl
Original file line number Diff line number Diff line change
@@ -48,7 +48,7 @@ export tensor, ⊗, compose

using Compat
using ChainRulesCore: ChainRulesCore, Tangent, ZeroTangent, NoTangent
using ChainRulesCore: @thunk, InplaceableThunk
using ChainRulesCore: @thunk, InplaceableThunk, ProjectTo, unthunk
using CompositionsBase
using Distances
using FillArrays
111 changes: 107 additions & 4 deletions src/chainrules.jl
Original file line number Diff line number Diff line change
@@ -112,15 +112,118 @@
function ChainRulesCore.rrule(s::Sinus, x::AbstractVector, y::AbstractVector)
d = x - y
sind = sinpi.(d)
abs2_sind_r = abs2.(sind) ./ s.r
val = sum(abs2_sind_r)
gradx = twoπ .* cospi.(d) .* sind ./ (s.r .^ 2)
abs2_sind_r = abs2.(sind) ./ s.r .^ 2
val = sum(abs2_sind_r)
simsurace marked this conversation as resolved.
Show resolved Hide resolved
gradx = twoπ .* cospi.(d) .* sind ./ s.r .^ 2

Check warning on line 117 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L115-L117

Added lines #L115 - L117 were not covered by tests
function evaluate_pullback::Any)
return (r=-2Δ .* abs2_sind_r,), Δ * gradx, -Δ * gradx
return (r=-2Δ .* abs2_sind_r ./ s.r,), Δ * gradx, -Δ * gradx

Check warning on line 119 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L119

Added line #L119 was not covered by tests
end
return val, evaluate_pullback
end

function ChainRulesCore.rrule(

Check warning on line 124 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L124

Added line #L124 was not covered by tests
::typeof(Distances.pairwise),
d::Sinus,
x::AbstractMatrix;
dims = 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
)
project_x = ProjectTo(x)
function pairwise_pullback(z̄)
Δ = unthunk(z̄)
n = size(x, dims)
= zero(x)
= zero(d.r)
if dims == 1
for j in 1:n, i in 1:n
xi = view(x, i, :)
xj = view(x, j, :)
ds = twoπ .* Δ[i, j] .* sinpi.(xi .- xj) .* cospi.(xi .- xj) ./ d.r .^ 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
.-= 2 .* Δ[i, j] .* sinpi.(xi .- xj) .^ 2 ./ d.r .^ 3
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assumes that rbar supports setindex! which is not true in general.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think in this case, it is because the field of Sinus is a Vector.

x̄[i, :] += ds
x̄[j, :] -= ds
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
elseif dims == 2
for j in 1:n, i in 1:n
xi = view(x, :, i)
xj = view(x, :, j)
ds = twoπ .* Δ[i, j] .* sinpi.(xi .- xj) .* cospi.(xi .- xj) ./ d.r .^ 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
.-= 2 .* Δ[i, j] .* sinpi.(xi .- xj) .^ 2 ./ d.r .^ 3
x̄[:, i] += ds
x̄[:, j] -= ds
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end

Check warning on line 153 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L130-L153

Added lines #L130 - L153 were not covered by tests
end
NoTangent(), (r=r̄,), @thunk(project_x(x̄))

Check warning on line 155 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L155

Added line #L155 was not covered by tests
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
return Distances.pairwise(d, x; dims), pairwise_pullback

Check warning on line 157 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L157

Added line #L157 was not covered by tests
end

function ChainRulesCore.rrule(

Check warning on line 160 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L160

Added line #L160 was not covered by tests
::typeof(Distances.pairwise),
d::Sinus,
x::AbstractMatrix,
y::AbstractMatrix;
dims = 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
)
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function pairwise_pullback(z̄)
Δ = unthunk(z̄)
n = size(x, dims)
m = size(y, dims)
= zero(x)
= zero(y)
= zero(d.r)
if dims == 1
for j in 1:m, i in 1:n
xi = view(x, i, :)
yj = view(y, j, :)
ds = twoπ .* Δ[i, j] .* sinpi.(xi .- yj) .* cospi.(xi .- yj) ./ d.r .^ 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
.-= 2 .* Δ[i, j] .* sinpi.(xi .- yj) .^ 2 ./ d.r .^ 3
x̄[i, :] += ds
ȳ[j, :] -= ds
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
elseif dims == 2
for j in 1:m, i in 1:n
xi = view(x, :, i)
yj = view(y, :, j)
ds = twoπ .* Δ[i, j] .* sinpi.(xi .- yj) .* cospi.(xi .- yj) ./ d.r .^ 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
.-= 2 .* Δ[i, j] .* sinpi.(xi .- yj) .^ 2 ./ d.r .^ 3
x̄[:, i] += ds
ȳ[:, j] -= ds
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end

Check warning on line 193 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L167-L193

Added lines #L167 - L193 were not covered by tests
end
NoTangent(), (r=r̄,), @thunk(project_x(x̄)), @thunk(project_y(ȳ))

Check warning on line 195 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L195

Added line #L195 was not covered by tests
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
return Distances.pairwise(d, x, y; dims), pairwise_pullback

Check warning on line 197 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L197

Added line #L197 was not covered by tests
end

function ChainRulesCore.rrule(

Check warning on line 200 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L200

Added line #L200 was not covered by tests
::typeof(Distances.colwise),
d::Sinus,
x::AbstractMatrix,
y::AbstractMatrix
simsurace marked this conversation as resolved.
Show resolved Hide resolved
)
project_x = ProjectTo(x)
project_y = ProjectTo(y)
function colwise_pullback(z̄)
Δ = unthunk(z̄)
n = size(x, 2)
= zero(x)
= zero(y)
= zero(d.r)
for i in 1:n
xi = view(x, :, i)
yi = view(y, :, i)
ds = twoπ .* Δ[i] .* sinpi.(xi .- yi) .* cospi.(xi .- yi) ./ d.r .^ 2
simsurace marked this conversation as resolved.
Show resolved Hide resolved
.-= 2 .* Δ[i] .* sinpi.(xi .- yi) .^ 2 ./ d.r .^ 3
x̄[:, i] += ds
ȳ[:, i] -= ds
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
NoTangent(), (r=r̄,), @thunk(project_x(x̄)), @thunk(project_y(ȳ))

Check warning on line 222 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L206-L222

Added lines #L206 - L222 were not covered by tests
simsurace marked this conversation as resolved.
Show resolved Hide resolved
end
return Distances.colwise(d, x, y), colwise_pullback

Check warning on line 224 in src/chainrules.jl

Codecov / codecov/patch

src/chainrules.jl#L224

Added line #L224 was not covered by tests
end

## Reverse Rules SqMahalanobis

function ChainRulesCore.rrule(
3 changes: 1 addition & 2 deletions test/basekernels/periodic.jl
Original file line number Diff line number Diff line change
@@ -15,7 +15,6 @@
TestUtils.test_interface(PeriodicKernel(; r=[0.9, 0.9]), ColVecs{Float64})
TestUtils.test_interface(PeriodicKernel(; r=[0.8, 0.7]), RowVecs{Float64})

# test_ADs(r->PeriodicKernel(r =exp.(r)), log.(r), ADs = [:ForwardDiff, :ReverseDiff])
@test_broken "Undefined adjoint for Sinus metric, and failing randomly for ForwardDiff and ReverseDiff"
test_ADs(r -> PeriodicKernel(r = exp.(r)), log.(r))
simsurace marked this conversation as resolved.
Show resolved Hide resolved
test_params(k, (r,))
end