Merge branch 'master' into bc/rm-boundscheck

Author: ToucheSir

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.65"
+version = "0.6.67"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -18,14 +18,24 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
-PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
+[weakdeps]
+Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+
+[extensions]
+ZygoteColorsExt = "Colors"
+ZygoteDistancesExt = "Distances"
+ZygoteTrackerExt = "Tracker"
+
 [compat]
 AbstractFFTs = "1.3.1"
 ChainRules = "1.44.1"
@@ -38,22 +48,18 @@ FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
 ForwardDiff = "0.10"
 GPUArrays = "8.4.2, 9"
 GPUArraysCore = "0.1.1"
-IRTools = "0.4.4"
+IRTools = "0.4.11"
 LogExpFunctions = "0.3.1"
 MacroTools = "0.5"
 NaNMath = "0.3, 1"
-Requires = "1.1"
 PrecompileTools = "1"
+Requires = "1.1"
 SpecialFunctions = "1.6, 2"
+Statistics = "1"
 Tracker = "0.2"
-ZygoteRules = "0.2.1"
+ZygoteRules = "0.2.4"
 julia = "1.6"
 
-[extensions]
-ZygoteColorsExt = "Colors"
-ZygoteDistancesExt = "Distances"
-ZygoteTrackerExt = "Tracker"
-
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
@@ -68,8 +74,3 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [targets]
 test = ["ChainRulesTestUtils", "Conda", "CUDA", "Distances", "FFTW", "FiniteDifferences", "PyCall", "Test"]
-
-[weakdeps]
-Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
-Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
-Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"

src/compiler/reverse.jl

@@ -248,7 +248,6 @@ Variable(a::Alpha) = Variable(a.id)
 sig(b::IRTools.Block) = unique([arg for br in branches(b) for arg in br.args if arg isa Variable])
 sig(pr::Primal) = Dict(b.id => sig(b) for b in blocks(pr.ir))
 
-# TODO unreachables?
 function adjointcfg(pr::Primal)
   ir = empty(pr.ir)
   return!(ir, nothing)
@@ -261,7 +260,9 @@ function adjointcfg(pr::Primal)
         push!(rb, xcall(Base, :(!==), alpha(pr.branches[b.id]), BranchNumber(i)))
       branch!(rb, preds[i].id, unless = cond)
     end
-    if !isempty(branches(b)) && branches(b)[end] == IRTools.unreachable
+    if isempty(preds) || (!isempty(branches(b)) && branches(b)[end] == IRTools.unreachable)
+      # If `b` is unreachable, then no context produced by the primal should end up branching to `rb`
+      push!(rb, xcall(Core, :throw, "unreachable")) # `throw` is necessary for inference not to hit the `unreachable`
       branch!(rb, 0)
     end
   end

src/deprecated.jl

@@ -65,3 +65,9 @@ macro nograd(ex)
   end
   return blk
 end
+
+# Internal function used by some downstream packages.
+# Removing this completely would require some tricky registry changes,
+# but leaving it as a vestigial function is much easier.
+# See https://github.com/FluxML/Zygote.jl/pull/1328 for more context.
+function ∇getindex end

src/lib/array.jl

@@ -41,24 +41,6 @@ end
 @adjoint (::Type{T})(sz) where {T<:Zeros} = T(sz), Δ->(nothing,)
 @adjoint (::Type{T})(sz) where {T<:Ones} = T(sz), Δ->(nothing,)
 
-@adjoint getindex(x::AbstractArray, inds...) = x[inds...], ∇getindex(x, inds)
-
-@adjoint view(x::AbstractArray, inds...) = view(x, inds...), ∇getindex(x, inds)
-
-∇getindex(x::AbstractArray{T,N}, inds) where {T,N} = dy -> begin
-  if inds isa NTuple{N,Int} && T <: Number
-    dx = OneElement(dy, inds, axes(x))
-  elseif inds isa NTuple{<:Any, Integer}
-    dx = _zero(x, typeof(dy))
-    dx[inds...] = dy
-  else
-    dx = _zero(x, eltype(dy))
-    dxv = view(dx, inds...)
-    dxv .= accum.(dxv, _droplike(dy, dxv))
-  end
-  return (_project(x, dx), map(_->nothing, inds)...)
-end
-
 """
     OneElement(val, ind, axes) <: AbstractArray
 

test/compiler.jl

@@ -227,10 +227,30 @@ end
 @test gradient(x -> sum(norm, collect(eachcol(x))), ones(3, 400))[1] ≈ fill(0.5773502691896258, 3, 400)
 
 # Tests adapted from https://github.com/dfdx/Umlaut.jl/pull/35
-@eval _boundscheck_foo(x) = ifelse($(Expr(:boundscheck)), 2x, x)
+@eval _has_boundscheck(x) = ifelse($(Expr(:boundscheck)), 2x, x)
 
 @testset "Meta Expr handling" begin
-  y, (dx,) = withgradient(_boundscheck_foo, 1)
+  y, (dx,) = withgradient(_has_boundscheck, 1)
   @test y == 2
   @test dx == 2
 end
+
+# issue 1118 & 1380
+function f_1380(x)
+    if rand(Bool)
+        return x
+    else
+        return 2x
+    end
+
+    # unreachable
+    return nothing
+end
+
+@testset "unreachable block" begin
+    y, back = Zygote.pullback(f_1380, 1.)
+    # There should not be a compiler error
+    local g
+    @test_nowarn g = back(1.)
+    @test only(g) ∈ (1., 2.)
+end

test/gradcheck.jl

@@ -174,11 +174,11 @@ end
 
   # Ensure that nothings work with numeric types.
   _, back = Zygote.pullback(getindex, randn(4), [1])
-  @test back([nothing]) == (zeros(4), nothing)
+  @test back([nothing]) === nothing
 
   # Ensure that nothings work with non-numeric types.
   _, back = Zygote.pullback(getindex, [randn(2) for _ in 1:3], [1])
-  @test back([nothing]) == (nothing, nothing)
+  @test back([nothing]) === nothing
 end
 
 @testset "view" begin

test/utils.jl

@@ -2,14 +2,8 @@ using ForwardDiff
 using Zygote: hessian_dual, hessian_reverse
 
 @testset "hessian: $hess" for hess in [hessian_dual, hessian_reverse]
-
-  if hess == hessian_dual
-    @test hess(x -> x[1]*x[2], randn(2)) ≈ [0 1; 1 0]
-    @test hess(((x,y),) -> x*y, randn(2)) ≈ [0 1; 1 0]  # original docstring version
-  else
-    @test_broken hess(x -> x[1]*x[2], randn(2)) ≈ [0 1; 1 0]  # can't differentiate ∇getindex
-    @test_broken hess(((x,y),) -> x*y, randn(2)) ≈ [0 1; 1 0]
-  end
+  @test hess(x -> x[1]*x[2], randn(2)) ≈ [0 1; 1 0]
+  @test hess(((x,y),) -> x*y, randn(2)) ≈ [0 1; 1 0]  # original docstring version
   @test hess(x -> sum(x.^3), [1 2; 3 4]) ≈ Diagonal([6, 18, 12, 24])
   @test hess(sin, pi/2) ≈ -1