Use handwritten rules for zero and one

Author: MasonProtter

src/rulesets/Base/base.jl

@@ -2,11 +2,38 @@
 # that also have FastMath versions.
 
 @scalar_rule copysign(y, x) (ifelse(signbit(x)!=signbit(y), -one(y), +one(y)), NoTangent())
-
-@scalar_rule one(x) ZeroTangent()
-@scalar_rule zero(x) ZeroTangent()
 @scalar_rule transpose(x) true
 
+# `zero`
+
+function frule((_, Δ1), ::typeof(zero), x)
+    var"∂f/∂x" = ZeroTangent()
+    (zero(x), Δ1 * var"∂f/∂x")
+end
+
+function rrule(::typeof(zero), x)
+    Ω = zero(x)
+    proj_x = ProjectTo(x)
+    var"∂f/∂x" = ZeroTangent()
+    pullback(Δ1) = (NoTangent(), proj_x(conj(var"∂f/∂x") * Δ1))
+    (Ω, pullback)
+end
+
+# `one`
+
+function frule((_, Δ1), ::typeof(one), x)
+    var"∂f/∂x" = ZeroTangent()
+    (one(x), Δ1 * var"∂f/∂x")
+end
+
+function rrule(::typeof(one), x)
+    Ω = one(x)
+    proj_x = ProjectTo(x)
+    var"∂f/∂x" = ZeroTangent()
+    pullback(Δ1) = (NoTangent(), proj_x(conj(var"∂f/∂x") * Δ1))
+    (Ω, pullback)
+end
+
 # `adjoint`
 
 frule((_, Δz), ::typeof(adjoint), z::Number) = (z', Δz')