Merge tag 'v0.6.59' into muse

[Diff since v0.6.58](FluxML/Zygote.jl@v0.6.58...v0.6.59)

**Merged pull requests:**
- Actually make sure conda env dir is set on Buildkite CI (FluxML#1392) (@ToucheSir)
- generated z2d (FluxML#1394) (@chengchingwen)
- bump version (FluxML#1398) (@chengchingwen)

Author: marius311

.buildkite/pipeline.yml

@@ -4,6 +4,8 @@ steps:
       - JuliaCI/julia#v1:
           version: "1.6"
       - JuliaCI/julia-test#v1: ~
+    command:
+      - mkdir -p "$${JULIA_DEPOT_PATH}/conda/3/x86_64"
     agents:
       queue: "juliagpu"
       cuda: "*"
@@ -14,6 +16,8 @@ steps:
       - JuliaCI/julia#v1:
           version: "1"
       - JuliaCI/julia-test#v1: ~
+    command:
+      - mkdir -p "$${JULIA_DEPOT_PATH}/conda/3/x86_64"
     agents:
       queue: "juliagpu"
       cuda: "*"

.github/workflows/Downstream.yml

@@ -27,14 +27,14 @@ jobs:
           - {user: SciML, repo: NeuralPDE.jl, group: NNPDE}
           - {user: JuliaMolSim, repo: Molly.jl, group: Zygote}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
       - uses: julia-actions/julia-buildpkg@latest
       - name: Clone Downstream
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           repository: ${{ matrix.package.user }}/${{ matrix.package.repo }}
           path: downstream

.github/workflows/TagBot.yml

@@ -4,6 +4,22 @@ on:
     types:
       - created
   workflow_dispatch:
+    inputs:
+      lookback:
+        default: 3
+permissions:
+  actions: read
+  checks: read
+  contents: write
+  deployments: read
+  issues: read
+  discussions: read
+  packages: read
+  pages: read
+  pull-requests: read
+  repository-projects: read
+  security-events: read
+  statuses: read
 jobs:
   TagBot:
     if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'

.github/workflows/ci.yml

@@ -32,12 +32,12 @@ jobs:
         #     version: '1'
         #     arch: x64
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v1
+      - uses: actions/cache@v3
         env:
           cache-name: cache-artifacts
         with:
@@ -58,15 +58,15 @@ jobs:
         #continue-on-error: ${{ matrix.version == 'nightly' }} # comment out to report nightly failures
       - uses: julia-actions/julia-processcoverage@v1
         if: matrix.version == '1' && matrix.os == 'ubuntu-latest'
-      - uses: codecov/codecov-action@v2
+      - uses: codecov/codecov-action@v3
         if: matrix.version == '1' && matrix.os == 'ubuntu-latest'
         with:
           file: lcov.info
   docs:
     name: Documentation
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1'

.github/workflows/clean_preview.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           ref: gh-pages
       - name: Delete preview and history + push changes

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.55"
+version = "0.6.59"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -20,17 +20,30 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 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 = "0.5, 1.0"
+AbstractFFTs = "1.3.1"
 ChainRules = "1.44.1"
 ChainRulesCore = "1.9"
 ChainRulesTestUtils = "1"
+Colors = "0.12"
 DiffRules = "1.4"
+Distances = "0.10"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13"
 ForwardDiff = "0.10"
 GPUArrays = "8.4.2"
@@ -40,17 +53,23 @@ LogExpFunctions = "0.3.1"
 MacroTools = "0.5"
 NaNMath = "0.3, 1"
 Requires = "1.1"
+SnoopPrecompile = "1.0.3"
 SpecialFunctions = "1.6, 2"
+Tracker = "0.2"
 ZygoteRules = "0.2.1"
 julia = "1.6"
 
 [extras]
+Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [targets]
-test = ["ChainRulesTestUtils", "CUDA", "Distances", "FFTW", "FiniteDifferences", "Test"]
+test = ["ChainRulesTestUtils", "Conda", "CUDA", "Distances", "FFTW", "FiniteDifferences", "PyCall", "Test"]

docs/src/utils.md

@@ -6,6 +6,7 @@ or a Hessian (by taking a second derivative).
 ```@docs
 Zygote.jacobian
 Zygote.hessian
+Zygote.hessian_reverse
 Zygote.diaghessian
 ```
 

ext/ZygoteColorsExt.jl

@@ -0,0 +1,13 @@
+module ZygoteColorsExt
+
+if isdefined(Base, :get_extension)
+    using Zygote
+    using Colors
+else
+    using ..Zygote
+    using ..Colors
+end
+
+Zygote.@non_differentiable Colors.ColorTypes._parameter_upper_bound(::Any...)
+
+end

src/lib/distances.jl renamed to ext/ZygoteDistancesExt.jl

@@ -1,4 +1,16 @@
-using .Distances
+module ZygoteDistancesExt
+
+if isdefined(Base, :get_extension)
+    using Zygote
+    using Distances
+    using LinearAlgebra
+else
+    using ..Zygote
+    using ..Distances
+    using ..LinearAlgebra
+end
+
+using Zygote: @adjoint, AContext, _pullback
 
 @adjoint function (::SqEuclidean)(x::AbstractVector, y::AbstractVector)
   δ = x .- y
@@ -66,22 +78,34 @@ end
 
 _sqrt_if_positive(d, δ) = d > δ ? sqrt(d) : zero(d)
 
-@adjoint function pairwise(dist::Euclidean, X::AbstractMatrix, Y::AbstractMatrix; dims=2)
+function Zygote._pullback(cx::AContext, ::Core.kwftype(typeof(pairwise)),
+                   kws::@NamedTuple{dims::Int}, ::typeof(pairwise), dist::Euclidean,
+                   X::AbstractMatrix, Y::AbstractMatrix)
   # Modify the forwards-pass slightly to ensure stability on the reverse.
+  dims = kws.dims
   function _pairwise_euclidean(sqdist::SqEuclidean, X, Y)
     D2 = pairwise(sqdist, X, Y; dims=dims)
     δ = eps(eltype(D2))
     return _sqrt_if_positive.(D2, δ)
   end
-  return pullback(_pairwise_euclidean, SqEuclidean(dist.thresh), X, Y)
+  res, back = _pullback(cx, _pairwise_euclidean, SqEuclidean(dist.thresh), X, Y)
+  pairwise_Euclidean_pullback(Δ) = (nothing, nothing, back(Zygote.unthunk_tangent(Δ))...)
+  return res, pairwise_Euclidean_pullback
 end
 
-@adjoint function pairwise(dist::Euclidean, X::AbstractMatrix; dims=2)
+function Zygote._pullback(cx::AContext, ::Core.kwftype(typeof(pairwise)),
+                   kws::@NamedTuple{dims::Int}, ::typeof(pairwise), dist::Euclidean,
+                   X::AbstractMatrix)
   # Modify the forwards-pass slightly to ensure stability on the reverse.
+  dims = kws.dims
   function _pairwise_euclidean(sqdist::SqEuclidean, X)
     D2 = pairwise(sqdist, X; dims=dims)
     δ = eps(eltype(D2))
     return _sqrt_if_positive.(D2, δ)
   end
-  return pullback(_pairwise_euclidean, SqEuclidean(dist.thresh), X)
+  res, back = _pullback(cx, _pairwise_euclidean, SqEuclidean(dist.thresh), X)
+  pairwise_Euclidean_pullback(Δ) = (nothing, nothing, back(Zygote.unthunk_tangent(Δ))...)
+  return res, pairwise_Euclidean_pullback
+end
+
 end

ext/ZygoteTrackerExt.jl

@@ -0,0 +1,17 @@
+module ZygoteTrackerExt
+
+if isdefined(Base, :get_extension)
+    using Zygote
+    using Tracker: Tracker, TrackedArray, TrackedReal
+else
+    using ..Zygote
+    using ..Tracker: Tracker, TrackedArray, TrackedReal
+end
+
+Zygote.unwrap(x::Union{TrackedArray,TrackedReal}) = Tracker.data(x)
+
+Zygote.pullback(f, ps::Tracker.Params) = pullback(f, ZygtParams(ps))
+Tracker.forward(f, ps::Params) = Tracker.forward(f, Tracker.Params(ps))
+Tracker.gradient_(f, ps::Params) = Tracker.gradient_(f, Tracker.Params(ps))
+
+end

src/Zygote.jl

@@ -4,7 +4,7 @@ using LinearAlgebra, Statistics
 using LinearAlgebra: copytri!, AbstractTriangular
 
 import ZygoteRules: @adjoint, @adjoint!, AContext, adjoint, _pullback, pullback,
-  literal_getproperty, literal_getfield
+  literal_getproperty, literal_getfield, unthunk_tangent
 
 using ChainRulesCore
 using ChainRules: ChainRules, rrule, unthunk, canonicalize
@@ -43,7 +43,6 @@ include("lib/forward.jl")
 include("lib/utils.jl")
 include("lib/range.jl")
 include("lib/logexpfunctions.jl")
-@init @require Distances="b4f34e82-e78d-54a5-968a-f98e89d6e8f7" include("lib/distances.jl")
 
 # we need to define this late, so that the genfuncs see lib.jl
 # Move using statements out of this file to help with sysimage building
@@ -53,12 +52,11 @@ include("compiler/interface2.jl")
 
 include("profiler/Profile.jl")
 
-@init @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
-  include("flux.jl")
-end
 
-@init @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" begin
-  @non_differentiable Colors.ColorTypes._parameter_upper_bound(::Any...)
+if !isdefined(Base, :get_extension)
+  @init @require Distances="b4f34e82-e78d-54a5-968a-f98e89d6e8f7" include("../ext/ZygoteDistancesExt.jl")
+  @init @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("../ext/ZygoteTrackerExt.jl")
+  @init @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" include("../ext/ZygoteColorsExt.jl")
 end
 
 using InteractiveUtils
@@ -79,8 +77,11 @@ macro profile(ex)
   end
 end
 
-## reverted due to https://github.com/SciML/DiffEqFlux.jl/issues/783
-# using SnoopPrecompile
-# @precompile_all_calls precompile()
+using SnoopPrecompile
+# This caused freezes on early 1.8 patch versions,
+# see https://github.com/SciML/DiffEqFlux.jl/issues/783
+@static if VERSION < v"1.8" || VERSION >= v"1.8.5"
+  @precompile_all_calls precompile()
+end
 
 end # module

src/compiler/chainrules.jl

@@ -312,16 +312,39 @@ end
 # Right now it uses a NamedTuple but not for fields of the AbstractDict struct
 z2d(dx::NamedTuple, primal::AbstractDict) = dx
 
-function z2d(delta::NamedTuple, primal::T) where T  # arbitrart struct
+function _z2d_struct_fallback(delta::NamedTuple, primal::T) where T
   fnames = fieldnames(T)
   deltas = map(n -> get(delta, n, nothing), fnames)
   primals = map(n -> getfield(primal, n), fnames)
   inner = map(z2d, deltas, primals)  # recurse into fields
-    if inner isa Tuple{Vararg{AbstractZero}}
+  if inner isa Tuple{Vararg{AbstractZero}}
     return NoTangent()  # collapse all-zero case
   else
     backing = NamedTuple{fnames}(inner)
-    return canonicalize(Tangent{T, typeof(backing)}(backing))
+    return Tangent{T, typeof(backing)}(backing)
+  end
+end
+
+function z2d(delta::NamedTuple, primal::T) where T  # arbitrart struct
+  if @generated
+    fnames = fieldnames(T)
+    N = length(fnames)
+    deltas = [ :($(Symbol(:delta_, fname)) = get(delta, $(QuoteNode(fname)), nothing)) for fname in fnames ]
+    primals = [ :($(Symbol(:primal_, fname)) = getfield(primal, $(QuoteNode(fname)))) for fname in fnames ]
+    inner = Expr(:tuple, [ :(z2d($(Symbol(:delta_, fname)), $(Symbol(:primal_, fname)))) for fname in fnames ]...)
+    return quote
+      $(deltas...)
+      $(primals...)
+      inner = $inner
+      if inner isa Tuple{Vararg{AbstractZero}}
+        return NoTangent()  # collapse all-zero case
+      else
+        backing = NamedTuple{$fnames}(inner)
+        return Tangent{T, typeof(backing)}(backing)
+      end
+    end
+  else
+    return _z2d_struct_fallback(delta, primal)
   end
 end