Skip to content

Avoid NaN-propagation in scalar rules #551

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.

Sign up for GitHub

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

Open
wants to merge 18 commits into
base: main
Choose a base branch
from
Open

Avoid NaN-propagation in scalar rules #551

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
ce76725
Make zeros hard in scalar rules
sethaxen Mar 13, 2022
d83965c
Test hard zeros in scalar rules
sethaxen Mar 13, 2022
fcc151f
Increment minor version number
sethaxen Mar 13, 2022
09c25fb
Merge branch 'main' into strongzeroscalar
sethaxen Oct 13, 2022
6eb7457
Apply suggestions from code review
sethaxen Oct 13, 2022
721877f
Make zero for NoTangent return NoTangent
sethaxen Oct 15, 2022
47c4009
Implement zero for NotImplemented
sethaxen Oct 15, 2022
b54eba3
Revert "Implement zero for NotImplemented"
sethaxen Oct 15, 2022
dc51058
Increment minor version number
sethaxen Oct 15, 2022
4431401
Apply suggestions from code review
sethaxen Oct 15, 2022
9983b71
Run formatter
sethaxen Oct 15, 2022
ddf1d7c
Merge branch 'main' into strongzeroscalar
sethaxen Dec 29, 2022
8852aa6
Add tests for not covered cases
sethaxen Dec 29, 2022
7f02ff7
Merge branch 'main' into strongzeroscalar
sethaxen Jan 18, 2023
06f8b67
Add strong_mul and strong_muladd
sethaxen Feb 4, 2023
e66a097
Fix random generation pre-v1.9
sethaxen Feb 4, 2023
b7bdee1
Add methods for NotImplemented
sethaxen Feb 5, 2023
3cdd108
Update tests
sethaxen Feb 5, 2023
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 Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "ChainRulesCore"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.13.0"
version = "1.14.0"

[deps]
Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
Expand Down
6 changes: 4 additions & 2 deletions src/rule_definition_tools.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,9 +295,11 @@ function propagation_expr(Δs, ∂s, _conj=false, proj=identity)
# Apply `muladd` iteratively.
# Explicit multiplication is only performed for the first pair of partial and gradient.
(∂s_1, Δs_1), _∂s_Δs_tail = Iterators.peel(zip(_∂s, Δs))
init_expr = :($∂s_1 * $Δs_1)
# zero gradients are treated as hard zeros. This avoids propagation of NaNs when
# partials are non-finite
init_expr = :(ifelse(iszero($Δs_1), zero($∂s_1), $∂s_1) * $Δs_1)
summed_∂_mul_Δs = foldl(_∂s_Δs_tail; init=init_expr) do ex, (∂s_i, Δs_i)
:(muladd($∂s_i, $Δs_i, $ex))
:(muladd(ifelse(iszero($Δs_i), zero($∂s_i), $∂s_i), $Δs_i, $ex))
end
return :($proj($summed_∂_mul_Δs))
end
Expand Down
23 changes: 23 additions & 0 deletions test/rule_definition_tools.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,29 @@ end
@test (NoTangent(), 0.0 - 1.0im) === rrule(make_imaginary, 2.0im)[2](1.0)
end

@testset "@scalar_rule strong zero (co)tangents" begin
suminv(x, y) = inv(x) + inv(y)
@scalar_rule suminv(x, y) (-(inv(x)^2), -(inv(y)^2))

@test frule((NoTangent(), 1.0, 1.0), suminv, 0.0, 1.0) === (Inf, -Inf)
@test frule((NoTangent(), ZeroTangent(), 1.0), suminv, 0.0, 1.0) === (Inf, -1.0)
@test frule((NoTangent(), 0.0, 1.0), suminv, 0.0, 1.0) === (Inf, -1.0)

@test frule((NoTangent(), 1.0, 1.0), suminv, 1.0, 0.0) === (Inf, -Inf)
@test frule((NoTangent(), 1.0, ZeroTangent()), suminv, 1.0, 0.0) === (Inf, -1.0)
@test frule((NoTangent(), 1.0, 0.0), suminv, 1.0, 0.0) === (Inf, -1.0)

@test rrule(suminv, 0.0, 1.0)[2](1.0) === (NoTangent(), -Inf, -1.0)
@test rrule(suminv, 0.0, 1.0)[2](ZeroTangent()) ===
(NoTangent(), ZeroTangent(), ZeroTangent())
@test rrule(suminv, 0.0, 1.0)[2](0.0) === (NoTangent(), 0.0, 0.0)

@test rrule(suminv, 1.0, 0.0)[2](1.0) === (NoTangent(), -1.0, -Inf)
@test rrule(suminv, 1.0, 0.0)[2](ZeroTangent()) ===
(NoTangent(), ZeroTangent(), ZeroTangent())
@test rrule(suminv, 1.0, 0.0)[2](0.0) === (NoTangent(), 0.0, 0.0)
end

@testset "Regression tests against #276 and #265" begin
# https://github.com/JuliaDiff/ChainRulesCore.jl/pull/276
# https://github.com/JuliaDiff/ChainRulesCore.jl/pull/265
Expand Down