un-collapse nothings in gradient

Author: ToucheSir

Project.toml

@@ -57,7 +57,7 @@ Requires = "1.1"
 SpecialFunctions = "1.6, 2"
 Statistics = "1"
 Tracker = "0.2"
-ZygoteRules = "0.2.4"
+ZygoteRules = "0.2.5"
 julia = "1.6"
 
 [extras]

src/compiler/interface.jl

@@ -67,6 +67,10 @@ sensitivity(y::Complex) = error("Output is complex, so the gradient is not defin
 sensitivity(y::AbstractArray) = error("Output is an array, so the gradient is not defined. Perhaps you wanted jacobian.")
 sensitivity(y) = error("Output should be scalar; gradients are not defined for output $(repr(y))")
 
+# Preserves output as tuple when gradients are collapsed
+_project_all(::NTuple{N}, ::Nothing) where {N} = ntuple(_ -> nothing, N)
+_project_all(x::Tuple, dx::Tuple) = map(_project, x, dx)
+
 """
     gradient(f, args...)
 
@@ -95,7 +99,7 @@ julia> gradient([7, 11], 0, 1) do x, y, d
 function gradient(f, args...)
   y, back = pullback(f, args...)
   grad = back(sensitivity(y))
-  isnothing(grad) ? nothing : map(_project, args, grad)
+  return _project_all(args, grad)
 end
 
 # Base.adjoint(f::Function) = x -> gradient(f, x)[1]  # piracy!
@@ -161,7 +165,7 @@ function withgradient(f, args...)
   else
     back(sensitivity(y))
   end
-  results = isnothing(grad) ? map(_ -> nothing, args) : map(_project, args, grad)
+  results = _project_all(args, grad)
   (val=y, grad=results)
 end
 
@@ -421,6 +425,9 @@ function pullback(f, ps::Params)
   end
 end
 
+# No conversion required here
+_project_all(_, dx::Grads) = dx
+
 # Code Reflection
 
 function code_ir(f, T)

test/lib/number.jl

@@ -3,8 +3,8 @@
     @test gradient(floor, 1) === (0.0,)
     @test gradient(ceil, 1) === (0.0,)
     @test gradient(round, 1) === (0.0,)
-    @test gradient(hash, 1) === nothing
-    @test gradient(div, 1, 2) === nothing
+    @test gradient(hash, 1) === (nothing,)
+    @test gradient(div, 1, 2) === (nothing, nothing)
   end
 
   @testset "basics" begin

test/structures.jl

@@ -64,5 +64,5 @@ end
   end
 
   m, b = Zygote._pullback(Zygote.Context(), nameof, M)
-  @test b(m) == (nothing, nothing)
+  @test b(m) === nothing
 end