Add ignore_derivatives (Zygote's dropgrad and ignore) (#229)

Author: gxyd

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.6.0"
+version = "1.7.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

docs/Manifest.toml

@@ -13,13 +13,13 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
 path = ".."
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.0.0-DEV"
+version = "1.7.0"
 
 [[Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
-git-tree-sha1 = "dc7dedc2c2aa9faf59a55c622760a25cbefbe941"
+git-tree-sha1 = "1a90210acd935f222ea19657f143004d2c2a1117"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.31.0"
+version = "3.38.0"
 
 [[Dates]]
 deps = ["Printf"]
@@ -73,9 +73,9 @@ version = "1.3.0"
 
 [[JSON]]
 deps = ["Dates", "Mmap", "Parsers", "Unicode"]
-git-tree-sha1 = "81690084b6198a2e1da36fcfda16eeca9f9f24e4"
+git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
 uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-version = "0.21.1"
+version = "0.21.2"
 
 [[LibCURL]]
 deps = ["LibCURL_jll", "MozillaCACerts_jll"]
@@ -122,9 +122,9 @@ uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
 
 [[Parsers]]
 deps = ["Dates"]
-git-tree-sha1 = "c8abc88faa3f7a3950832ac5d6e690881590d6dc"
+git-tree-sha1 = "9d8c00ef7a8d110787ff6f170579846f776133a9"
 uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
-version = "1.1.0"
+version = "2.0.4"
 
 [[Pkg]]
 deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]

docs/make.jl

@@ -56,6 +56,7 @@ makedocs(
         "Gradient Accumulation" => "gradient_accumulation.md",
         "Usage in AD" => "use_in_ad_system.md",
         "Converting ZygoteRules" => "converting_zygoterules.md",
+        "Tips for making packages work with AD" => "tips_for_packages.md",
         "Design" => [
             "Changing the Primal" => "design/changing_the_primal.md",
             "Many Differential Types" => "design/many_differentials.md",

docs/src/api.md

@@ -46,6 +46,12 @@ Private = false
 ProjectTo
 ```
 
+## Ignoring gradients
+```@docs
+ignore_derivatives
+@ignore_derivatives
+```
+
 ## Internal
 ```@docs
 ChainRulesCore.AbstractTangent

docs/src/tips_for_packages.md

@@ -0,0 +1,46 @@
+# Tips for making your package work with AD
+
+## Ignoring gradients for certain expressions
+
+There exists code that is not meant to be differentiated through, for example logging.
+In some cases, AD systems might work perfectly well with that code, but in others they might not.
+A convenience function `ignore_derivatives` is provided to get around this issue.
+It captures the functionality of both `Zygote.ignore` and `Zygote.dropgrad`.
+
+For example, Zygote does not support mutation, so it will break if you try to store intermediate values as in the following example:
+```julia
+somes = []
+things = []
+
+function loss(x, y)
+    some = f(x, y)
+    thing = g(x)
+    
+    # log
+    push!(somes, some)
+    push!(things, thing)
+
+    return some + thing
+end
+```
+
+It is possible to get around this by using the `ignore_derivatives` function
+```julia
+ignore_derivatives() do
+    push!(somes, some)
+    push!(things, thing)
+end
+```
+or using a macro for one-liners
+```julia
+@ignore_derivatives push!(things, thing)
+```
+
+It is also possible to use this on individual objects, e.g.
+```julia
+ignore_derivatives(a) + b
+```
+will ignore the gradients for `a` only.
+
+Passing in instances of functors (callable structs), `ignore_derivatives(functor)`, will make them behave like normal structs, i.e. propagate without being called and dropping their gradients.
+If you want to call a functor in the primal computation, wrap it in a closure: `ignore_derivatives(() -> functor())`

src/ChainRulesCore.jl

@@ -13,6 +13,7 @@ export frule_via_ad, rrule_via_ad
 export @non_differentiable, @opt_out, @scalar_rule, @thunk, @not_implemented
 export ProjectTo, canonicalize, unthunk  # differential operations
 export add!!  # gradient accumulation operations
+export ignore_derivatives, @ignore_derivatives
 # differentials
 export Tangent, NoTangent, InplaceableThunk, Thunk, ZeroTangent, AbstractZero, AbstractThunk
 
@@ -32,6 +33,7 @@ include("projection.jl")
 include("config.jl")
 include("rules.jl")
 include("rule_definition_tools.jl")
+include("ignore_derivatives.jl")
 
 include("deprecated.jl")
 

src/ignore_derivatives.jl

@@ -0,0 +1,51 @@
+"""
+    ignore_derivatives(f::Function)
+
+Tells the AD system to ignore the gradients of the wrapped closure. The primal computation
+(forward pass) is executed normally.
+
+```julia
+ignore_derivatives() do
+    value = rand()
+    push!(collection, value)
+end
+```
+
+Using this incorrectly could lead to incorrect gradients.
+For example, the following function will have zero gradients with respect to its argument:
+```julia
+function wrong_grads(x)
+    y = ones(3)
+    ignore_derivatives() do
+        push!(y, x)
+    end
+    return sum(y)
+end
+```
+"""
+ignore_derivatives(f::Function) = f()
+
+"""
+    ignore_derivatives(x)
+
+Tells the AD system to ignore the gradients of the argument. Can be used to avoid
+unnecessary computation of gradients.
+
+```julia
+ignore_derivatives(x) * w
+```
+"""
+ignore_derivatives(x) = x
+
+@non_differentiable ignore_derivatives(f)
+
+"""
+    @ignore_derivatives (...)
+
+Tells the AD system to ignore the expression. Equivalent to `ignore_derivatives() do (...) end`.
+"""
+macro ignore_derivatives(ex)
+    return :(ChainRulesCore.ignore_derivatives() do
+        $(esc(ex))
+    end)
+end

test/ignore_derivatives.jl

@@ -0,0 +1,52 @@
+struct MyFunctor
+    a::Float64
+end
+(mf::MyFunctor)(b) = mf.a + b
+
+@testset "ignore_derivatives.jl" begin
+    @testset "function" begin
+        f() = return 4.0
+
+        y, ẏ = frule((1.0, ), ignore_derivatives, f)
+        @test y == f()
+        @test ẏ == NoTangent()
+
+        y, pb = rrule(ignore_derivatives, f)
+        @test y == f()
+        @test pb(1.0) == (NoTangent(), NoTangent())
+    end
+
+    @testset "argument" begin
+        arg = 2.1
+
+        y, ẏ = frule((1.0, ), ignore_derivatives, arg)
+        @test y == arg
+        @test ẏ == NoTangent()
+
+        y, pb = rrule(ignore_derivatives, arg)
+        @test y == arg
+        @test pb(1.0) == (NoTangent(), NoTangent())
+    end
+
+    @testset "functor" begin
+        mf = MyFunctor(1.0)
+
+        # as an argument
+        y, ẏ = frule((1.0,), ignore_derivatives, mf)
+        @test y == mf
+        @test ẏ == NoTangent()
+
+        y, pb = rrule(ignore_derivatives, mf)
+        @test y == mf
+        @test pb(1.0) == (NoTangent(), NoTangent())
+
+        # when called
+        y, ẏ = frule((1.0,), ignore_derivatives, ()->mf(3.0))
+        @test y == mf(3.0)
+        @test ẏ == NoTangent()
+
+        y, pb = rrule(ignore_derivatives, ()->mf(3.0))
+        @test y == mf(3.0)
+        @test pb(1.0) == (NoTangent(), NoTangent())
+    end
+end

test/runtests.jl

@@ -21,6 +21,7 @@ using Test
     include("rules.jl")
     include("rule_definition_tools.jl")
     include("config.jl")
+    include("ignore_derivatives.jl")
 
     include("deprecated.jl")
 end