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
sethaxen 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 all 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.15.7"
version = "1.16.0"

[deps]
Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
Expand Down
31 changes: 28 additions & 3 deletions src/rule_definition_tools.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,12 +292,14 @@ function propagation_expr(Δs, ∂s, _conj=false, proj=identity)
end
end

# Apply `muladd` iteratively.
# Apply `strong_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 = :(strong_mul($∂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))
:(strong_muladd($∂s_i, $Δs_i, $ex))
end
return :($proj($summed_∂_mul_Δs))
end
Expand Down Expand Up @@ -615,3 +617,26 @@ function _constrain_and_name(arg::Expr, _)
return error("malformed arguments: $arg")
end
_constrain_and_name(name::Symbol, constraint) = Expr(:(::), name, constraint) # add type

"""
strong_mul(x, y)

Multiply `x` and `y`. If `iszero(y)`, treat `y` as a hard zero even for non-finite `x`.
"""
strong_mul(x, y) = ifelse(iszero(y), zero(x), x) * y

@devmotion devmotion Oct 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Based on the docstring, I'd expect something like

Suggested change
strong_mul(x, y) = ifelse(iszero(y), zero(x), x) * y
function strong_mul(x, y)
xy = x * y
return iszero(y) ? zero(xy) : xy
end

ie adjusting the product i stead of x when treating y as a strong zero.

Any particular (performance?) reason for the current code?


"""
strong_muladd(x, y, z)

Multiply `x` and `y` and add to `z`. If `iszero(y)`, treat `y` as a hard zero even for
non-finite `x`.
"""
strong_muladd(x, y, z) = muladd(ifelse(iszero(y), zero(x), x), y, z)

@devmotion devmotion Oct 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Similarly, instead of zeroing x I would have expected zeroing the result of the product in the presence of a strong zero y, something like

Suggested change
strong_muladd(x, y, z) = muladd(ifelse(iszero(y), zero(x), x), y, z)
function strong_muladd(x, y, z)
xyz = muladd(x, y, z)
return iszero(y) ? of type(xyz, z) : xyz
end

I don't think it would matter, I'm mainly curious why it's not using the arguably more direct translation of the docstring.


# slightly faster for BigFloats
strong_mul(x::BigFloat, y::BigFloat) = (iszero(y) ? zero(x) : x) * y
strong_muladd(x::BigFloat, y::BigFloat, z) = muladd((iszero(y) ? zero(x) : x), y, z)

# avoid raising errors for NotImplemented
strong_mul(x::NotImplemented, y) = (iszero(y) ? zero(x) : x) * y
strong_muladd(x::NotImplemented, y, z) = muladd((iszero(y) ? zero(x) : x), y, z)
3 changes: 3 additions & 0 deletions src/tangent_types/abstract_zero.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,6 @@ arguments.
```
"""
struct NoTangent <: AbstractZero end

Base.zero(::NoTangent) = NoTangent()

@devmotion devmotion Oct 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

These seem reasonable, but can you remind me: why are these changes in the PR here? I'd expect to mainly see definitions and uses of strong_mul and strong_muladd in this PR.

Base.zero(::Type{NoTangent}) = NoTangent()
109 changes: 109 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 @@ -224,6 +224,61 @@ end
end
end

@testset "strong_mul" begin
ni = @not_implemented("not implemented!")
xvals = (
5,
randn(Float32),
randn(Float64),
randn(ComplexF64),
big(randn()),
ZeroTangent(),
)
yvals = (3, randn(Float32), randn(Float64), randn(ComplexF64), big(randn()))
yzerovals = (0, 0.0f0, 0.0, 0.0im, big(0.0), ZeroTangent(), NoTangent())
@testset for x in xvals
x === ni || @testset for y in yvals
@test @inferred(ChainRulesCore.strong_mul(x, y)) == x * y
end
@testset for y in yzerovals
@test @inferred(ChainRulesCore.strong_mul(x, y)) == zero(x * y)
if x isa AbstractFloat
@test ChainRulesCore.strong_mul(oftype(x, Inf), y) == zero(x * y)
@test ChainRulesCore.strong_mul(oftype(x, -Inf), y) == zero(x * y)
@test ChainRulesCore.strong_mul(oftype(x, NaN), y) == zero(x * y)
end
end
end
end

@testset "strong_muladd" begin
ni = @not_implemented("not implemented!")
xvals = (
5,
randn(Float32),
randn(Float64),
randn(ComplexF64),
big(randn()),
ZeroTangent(),
)
yvals = (3, randn(Float32), randn(Float64), randn(ComplexF64), big(randn()))
zvals = (7, randn(Float32), randn(Float64), randn(ComplexF64), big(randn()))
yzerovals = (0, 0.0f0, 0.0, 0.0 * im, big(0.0), ZeroTangent(), NoTangent())
@testset for x in xvals, z in zvals
x === ni || @testset for y in yvals
@test @inferred(ChainRulesCore.strong_muladd(x, y, z)) == muladd(x, y, z)
end
@testset for y in yzerovals
@test @inferred(ChainRulesCore.strong_muladd(x, y, z)) == z
if x isa AbstractFloat
@test ChainRulesCore.strong_muladd(oftype(x, Inf), y, z) == z
@test ChainRulesCore.strong_muladd(oftype(x, -Inf), y, z) == z
@test ChainRulesCore.strong_muladd(oftype(x, NaN), y, z) == z
end
end
end
end

@testset "@scalar_rule" begin
@testset "@scalar_rule with multiple output" begin
simo(x) = (x, 2x)
Expand Down Expand Up @@ -256,6 +311,60 @@ 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 @inferred(frule((NoTangent(), 1.0, 1.0), suminv, 0.0, 1.0)) ===
(Inf, -Inf)
@test @inferred(frule((NoTangent(), ZeroTangent(), 1.0), suminv, 0.0, 1.0)) ===
(Inf, -1.0)
@test @inferred(frule((NoTangent(), NoTangent(), 1.0), suminv, 0.0, 1.0)) ===
(Inf, -1.0)
@test @inferred(frule((NoTangent(), 0.0, 1.0), suminv, 0.0, 1.0)) ===
(Inf, -1.0)

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

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

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

# cases not covered
t = @thunk(0.0)
@inferred(frule((NoTangent(), t, 1.0), suminv, 0.0, 1.0))
@inferred(frule((NoTangent(), 1.0, t), suminv, 1.0, 0.0))
@inferred(rrule(suminv, 0.0, 1.0)[2](t))
@inferred(rrule(suminv, 1.0, 0.0)[2](t))
@test_broken rrule(suminv, 0.0, 1.0)[2](t) == (NoTangent(), 0.0, 0.0)
@test_broken rrule(suminv, 1.0, 0.0)[2](t) == (NoTangent(), 0.0, 0.0)
@test_broken frule((NoTangent(), t, 1.0), suminv, 0.0, 1.0) == (Inf, -1.0)
@test_broken frule((NoTangent(), 1.0, t), suminv, 1.0, 0.0) == (Inf, -1.0)

ni = @not_implemented("not implemented!")
@test_broken rrule(suminv, 0.0, 1.0)[2](ni) == (NoTangent(), 0.0, 0.0)
@test_broken rrule(suminv, 1.0, 0.0)[2](ni) == (NoTangent(), 0.0, 0.0)
@test_broken frule((NoTangent(), ni, 1.0), suminv, 0.0, 1.0) == (Inf, -1.0)
@test_broken frule((NoTangent(), 1.0, ni), suminv, 1.0, 0.0) == (Inf, -1.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
4 changes: 2 additions & 2 deletions test/tangent_types/abstract_zero.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,7 @@
end
@test broadcastable(z) isa Ref{ZeroTangent}
@test zero(@thunk(3)) === z
@test zero(NoTangent()) === z
@test zero(ZeroTangent) === z
@test zero(NoTangent) === z
@test zero(Tangent{Tuple{Int,Int}}((1, 2))) === z
for f in (transpose, adjoint, conj)
@test f(z) === z
Expand Down Expand Up @@ -94,6 +92,8 @@

@testset "NoTangent" begin
dne = NoTangent()
@test zero(dne) === NoTangent()
@test zero(NoTangent) === NoTangent()
@test dne + dne == dne
@test dne + 1 == 1
@test 1 + dne == 1
Expand Down