Merge pull request #1249 from mzgubic/mz/deprecate_nograd

deprecate `Zygote.@nograd`

Project.toml

@@ -27,7 +27,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 AbstractFFTs = "0.5, 1.0"
-ChainRules = "1.36.2"
+ChainRules = "1.37"
 ChainRulesCore = "1.9"
 ChainRulesTestUtils = "1"
 DiffRules = "1.4"

src/Zygote.jl

@@ -58,7 +58,7 @@ include("profiler/Profile.jl")
 end
 
 @init @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" begin
-  @nograd Colors.ColorTypes._parameter_upper_bound
+  @non_differentiable Colors.ColorTypes._parameter_upper_bound(::Any...)
 end
 
 using InteractiveUtils

src/deprecated.jl

@@ -49,3 +49,19 @@ macro ignore(ex)
         $(esc(ex))
     end)
 end
+
+using MacroTools: @q
+
+macro nograd(ex)
+  Base.depwarn(
+    "`Zygote.@nograd myfunc` is deprecated, use `ChainRulesCore.@non_differentiable myfunc(::Any...)` instead.",
+    :nograd
+  )
+  isexpr(ex, :tuple) || (ex = Expr(:tuple, ex))
+  blk = @q begin end
+  for f in ex.args
+    back = MacroTools.@q _ -> ($__source__; nothing)
+    push!(blk.args, :(@inline Zygote._pullback(::Context, ::Core.Typeof($(esc(f))), args...) = $(esc(f))(args...), $back))
+  end
+  return blk
+end

src/forward/lib.jl

@@ -9,7 +9,7 @@ end
 # TODO figure out why this made a test fail
 zerolike(x::Union{Module,Type}) = nothing
 
-# TODO: `@nograd` and `@linear`
+# TODO: `@non_differentiable` and `@linear`
 
 @tangent zerolike(x) = zerolike(x), _ -> zerolike(x)
 @tangent one(x::Number) = one(x), _ -> zero(x)

src/lib/array.jl

@@ -6,8 +6,6 @@ using Distributed: pmap, AbstractWorkerPool
 @adjoint Array(xs::AbstractArray) = Array(xs), ȳ -> (ȳ,)
 @adjoint Array(xs::Array) = Array(xs), ȳ -> (ȳ,)
 
-@nograd ones, zeros, Base.OneTo, Colon(), one, zero, sizehint!, count
-
 @adjoint copy(x::AbstractArray) = copy(x), ȳ -> (ȳ,)
 
 @adjoint collect(x::Tuple) = collect(x), dy -> (Tuple(dy),)
@@ -233,11 +231,6 @@ end
   end
 end
 
-for t in subtypes(AbstractWorkerPool)
-  @nograd t
-end
-@nograd workers
-
 function _pullback(cx::AContext, ::typeof(collect), g::Base.Generator)
   y, b = ∇map(cx, g.f, g.iter)
   back(::Nothing) = nothing

src/lib/base.jl

@@ -49,8 +49,6 @@ end
 
 # Channels
 
-@nograd Channel
-
 grad_mut(ch::Channel) = Channel(ch.sz_max)
 
 @adjoint! function put!(ch::Channel, x)
@@ -157,8 +155,6 @@ end
 
 @adjoint Base.nameof(x::UnionAll) = nameof(x), _ -> (nothing,)
 
-@nograd typeintersect
-
 # Base.Fix1 and Base.Fix2: https://github.com/FluxML/Zygote.jl/issues/957
 @adjoint function (g::Base.Fix1)(y)
     f = g.f

src/lib/buffer.jl

@@ -1,7 +1,7 @@
 grad_mut(b::Buffer) = fill!(similar(b.data, Any), nothing)
 grad_mut(b::Buffer{T}) where T<:Number = fill!(similar(b.data, float(T)), 0)
 
-@nograd Buffer
+@non_differentiable Buffer(::Any...)
 
 @adjoint function getindex(b::Buffer, i...)
   b[i...], function (Δ)

src/lib/grad.jl

@@ -1,15 +1,3 @@
-using MacroTools: @q
-
-macro nograd(ex)
-  isexpr(ex, :tuple) || (ex = Expr(:tuple, ex))
-  blk = @q begin end
-  for f in ex.args
-    back = MacroTools.@q _ -> ($__source__; nothing)
-    push!(blk.args, :(@inline Zygote._pullback(::Context, ::Core.Typeof($(esc(f))), args...) = $(esc(f))(args...), $back))
-  end
-  return blk
-end
-
 macro which(ex)
   @capture(ex, f_(args__)) || error("Zygote.@which f(args...)")
   :(InteractiveUtils.@which adjoint(Context(), $(esc(f)), $(esc.(args)...)))

src/lib/lib.jl

@@ -38,8 +38,6 @@ function accum(x::RefValue, y::RefValue)
 end
 
 # Core functions
-@nograd eps, Base.eval, Core.TypeVar, Core.UnionAll, Symbol
-
 @adjoint deepcopy(x) = deepcopy(x), ȳ -> (ȳ,)
 
 @adjoint (::Type{V})(x...) where V<:Val = V(x...), _ -> nothing

src/lib/number.jl

@@ -1,6 +1,3 @@
-
-@nograd floor, ceil, trunc, round, div
-
 @adjoint Base.literal_pow(::typeof(^), x::Number, ::Val{p}) where {p} =
   Base.literal_pow(^,x,Val(p)),
   Δ -> (nothing, Δ * conj(p * Base.literal_pow(^,x,Val(p-1))), nothing)

test/lib/number.jl

@@ -1,7 +1,7 @@
 @testset "nograds" begin
-  @test gradient(floor, 1) === nothing
-  @test gradient(ceil, 1) === nothing
-  @test gradient(round, 1) === nothing
+  @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
 end #testset