simplify, some doctests

Author: mcabbott

Project.toml

@@ -6,7 +6,6 @@ version = "0.6.20"
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -27,7 +26,6 @@ AbstractFFTs = "0.5, 1.0"
 ChainRules = "1.5"
 ChainRulesCore = "1.3"
 ChainRulesTestUtils = "1"
-Compat = "2.2, 3"
 DiffRules = "1.0"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12"
 ForwardDiff = "0.10"

README.md

@@ -18,7 +18,7 @@ julia> using Zygote
 julia> f(x) = 5x + 3
 
 julia> f(10), f'(10)
-(53, 5)
+(53, 5.0)
 
 julia> @code_llvm f'(10)
 define i64 @"julia_#625_38792"(i64) {

src/Zygote.jl

@@ -11,7 +11,6 @@ using ChainRules: ChainRules, rrule, unthunk, canonicalize
 using IRTools
 using MacroTools, Requires
 using MacroTools: @forward
-using Compat  # for Julia 1.3, need Compat 2.2
 
 import Distributed: pmap, CachingPool, workers
 export Params, withgradient, gradient, withjacobian, jacobian, hessian, diaghessian, pullback, pushforward, @code_adjoint

src/compiler/interface.jl

@@ -68,14 +68,13 @@ julia> gradient([7, 11], 0, 1) do x, y, d
          p = size(x, d)
          sum(x.^p .+ y)
        end
-([14.0, 22.0], 2, nothing)
+([14.0, 22.0], 2.0, nothing)
 ```
 """
 function gradient(f, args...)
   y, back = pullback(f, args...)
   grad = back(sensitivity(y))
-  isnothing(grad) && return nothing
-  map(_project, args, grad)
+  isnothing(grad) ? nothing : map(_project, args, grad)
 end
 
 Base.adjoint(f::Function) = x -> gradient(f, x)[1]
@@ -98,8 +97,8 @@ true
 function withgradient(f, args...)
   y, back = pullback(f, args...)
   grad = back(sensitivity(y))
-  isnothing(grad) && return (val=y, grad=nothing)
-  (val = y, grad = map(_project, args, grad))
+  results = isnothing(grad) ? map(_ -> nothing, args) : map(_project, args, grad)
+  (val=y, grad=results)
 end
 
 # Param-style wrappers
@@ -138,23 +137,17 @@ julia> g = gradient(Params([x, y])) do
 Grads(...)
 
 julia> g[x]
-2×3 Matrix{Int64}:
- 7  70  700
- 8  80  800
+2×3 Matrix{Float64}:
+ 7.0  70.0  700.0
+ 8.0  80.0  800.0
 
 julia> haskey(g, z)  # only x and y are parameters
 false
 ```
 """
-function gradient(f, ps::Params)
-  y, back = pullback(f, ps)
-  back(sensitivity(y))
-end
+gradient
 
-function withgradient(f, ps::Params)
-  y, back = pullback(f, ps)
-  (val = y, grad = back(sensitivity(y)))
-end
+Base.map(::typeof(_project), args::Tuple{Params}, grad) = grad  # skip _project in gradient(f, ::Params)
 
 function Base.union!(ps::Params, itrs...)
   foreach(itr -> foreach(x -> push!(ps, x), itr), itrs)

src/lib/broadcast.jl

@@ -51,7 +51,7 @@ function unbroadcast(x::AbstractArray, x̄)
     _project(x, x̄)  # ProjectTo handles reshape, offsets, structured matrices, row vectors
   else
     tup = filter(d -> size(x, d) == 1, ntuple(identity, N))
-    dims = length(tup) == 1 ? only(tup) : tup  # avoid sum(xbar, dims=(1,)) as e.g. sum(SA[1 2; 3 4], dims=(1,)) fails
+    dims = length(tup) == 1 ? first(tup) : tup  # avoid sum(xbar, dims=(1,)) as e.g. sum(SA[1 2; 3 4], dims=(1,)) fails
     _project(x, accum_sum(x̄; dims = dims))
   end
 end