muladd with Zero() (#98)

* `muladd` with `Zero()`

Co-authored-by: "Shashi Gowda" <gowda@mit.edu>
Co-authored-by: "Yingbo Ma" <mayingbo5@gmail.com>

* Bump version

Co-authored-by: Shashi Gowda <shashigowda91@gmail.com>

Author: YingboMa

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.5.4"
+version = "0.5.4-DEV"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

src/differential_arithmetic.jl

@@ -31,6 +31,15 @@ Base.:+(::Zero, ::DoesNotExist) = DoesNotExist()
 Base.:*(::DoesNotExist, ::Zero) = Zero()
 Base.:*(::Zero, ::DoesNotExist) = Zero()
 
+Base.muladd(::Zero, x, y) = y
+Base.muladd(x, ::Zero, y) = y
+Base.muladd(x, y, ::Zero) = x*y
+
+Base.muladd(::Zero, ::Zero, y) = y
+Base.muladd(x, ::Zero, ::Zero) = Zero()
+Base.muladd(::Zero, x, ::Zero) = Zero()
+
+Base.muladd(::Zero, ::Zero, ::Zero) = Zero()
 
 Base.:+(::Zero, ::Zero) = Zero()
 Base.:*(::Zero, ::Zero) = Zero()

test/differentials/zero.jl

@@ -16,4 +16,14 @@
     @test zero(One()) === z
     @test zero(DoesNotExist()) === z
     @test zero(Composite{Tuple{Int,Int}}((1, 2))) === z
+
+    # use mutable objects to test the strong `===` condition
+    x = ones(2)
+    @test muladd(Zero(), 2, x) === x
+    @test muladd(2, Zero(), x) === x
+    @test muladd(Zero(), Zero(), x) === x
+    @test muladd(2, 2, Zero()) === 4
+    @test muladd(x, Zero(), Zero()) === Zero()
+    @test muladd(Zero(), x, Zero()) === Zero()
+    @test muladd(Zero(), Zero(), Zero()) === Zero()
 end