Add iszero(x) branches to xlogy and xlog1py (#57)

* Add `iszero(x)` branches to `xlogy` and `xlog1py`

* Import FiniteDifferences for tests

* Remove additional NaN branch

* Fix tests

* White space

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

* White space

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

* Simplify expression

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

* Improve formatting

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

* Adjust and fix tests

* Correct mistake

* Revert `frule` change, update test accordingly

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

Author: simsurace

Project.toml

@@ -21,9 +21,10 @@ julia = "1"
 
 [extras]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesTestUtils", "OffsetArrays", "Random", "Test"]
+test = ["ChainRulesTestUtils", "FiniteDifferences", "OffsetArrays", "Random", "Test"]

src/chainrules.jl

@@ -19,9 +19,15 @@ end
 function _Ω_∂_xlogy(x::Real, y::Real)
     logy = log(y)
     z = x * logy
-    Ω = iszero(x) && !isnan(y) ? zero(z) : z
+    w = x / y
+    if iszero(x) && !isnan(y)
+        Ω = zero(z)
+        ∂y = zero(w)
+    else
+        Ω = z
+        ∂y = w
+    end
     ∂x = logy
-    ∂y = x / y
     return Ω, ∂x, ∂y
 end
 function ChainRulesCore.frule((_, Δx, Δy), ::typeof(xlogy), x::Real, y::Real)
@@ -38,9 +44,15 @@ end
 function _Ω_∂_xlog1py(x::Real, y::Real)
     log1py = log1p(y)
     z = x * log1py
-    Ω = iszero(x) && !isnan(y) ? zero(z) : z
+    w = x / (1 + y)
+    if iszero(x) && !isnan(y)
+        Ω = zero(z)
+        ∂y = zero(w)
+    else
+        Ω = z
+        ∂y = w
+    end
     ∂x = log1py
-    ∂y = x / (1 + y)
     return Ω, ∂x, ∂y
 end
 function ChainRulesCore.frule((_, Δx, Δy), ::typeof(xlog1py), x::Real, y::Real)

test/chainrules.jl

@@ -2,6 +2,13 @@
     x = rand()
     test_frule(xlogx, x)
     test_rrule(xlogx, x)
+
+    # Test `iszero(x)` branches
+    test_frule(xlogy, 0.0, 1.0; fdm = forward_fdm(5, 1), nans = true)
+    test_rrule(xlogy, 0.0, 1.0; fdm = forward_fdm(5, 1), nans = true)
+    @test iszero(last(frule((NoTangent(), ZeroTangent(), 1.), xlog1py, 0.0, -1.0)))
+    @test iszero(last(last(rrule(xlog1py, 0.0, -1.0))(1.)))
+
     for x in (-x, 0.0, x)
         y = rand()
         test_frule(xlogy, x, y)

test/runtests.jl

@@ -2,6 +2,7 @@ using LogExpFunctions
 using ChainRulesTestUtils
 using ChainRulesCore
 using ChangesOfVariables
+using FiniteDifferences
 using InverseFunctions
 using OffsetArrays