Merge branch 'lecture_09'

pevnak · pevnak · commit 12fdfd894ef7 · 2021-11-25T21:45:00.000+01:00
diff --git a/docs/src/lecture_08/Manifest.toml b/docs/src/lecture_08/Manifest.toml
@@ -473,6 +473,7 @@ version = "1.3.5"
 deps = ["Libdl"]
 uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
+
 [[LittleCMS_jll]]
 deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "Pkg"]
 git-tree-sha1 = "110897e7db2d6836be22c18bffd9422218ee6284"
@@ -551,6 +552,7 @@ git-tree-sha1 = "7937eda4681660b4d6aeeecc2f7e1c81c8ee4e2f"
 uuid = "e7412a2a-1a6e-54c0-be00-318e2571c051"
 version = "1.3.5+0"
 
+
 [[OpenJpeg_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "Pkg", "libpng_jll"]
 git-tree-sha1 = "76374b6e7f632c130e78100b166e5a48464256f8"
@@ -624,6 +626,7 @@ git-tree-sha1 = "0d185e8c33401084cab546a756b387b15f76720c"
 uuid = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 version = "1.23.6"
 
+
 [[Poppler_jll]]
 deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "Glib_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "Pkg", "libpng_jll"]
 git-tree-sha1 = "e11443687ac151ac6ef6699eb75f964bed8e1faa"
@@ -688,6 +691,7 @@ git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
 uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.1.3"
 
+
 [[Richardson]]
 deps = ["LinearAlgebra"]
 git-tree-sha1 = "e03ca566bec93f8a3aeb059c8ef102f268a38949"
@@ -788,6 +792,7 @@ version = "1.6.0"
 deps = ["ArgTools", "SHA"]
 uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 
+
 [[Tectonic]]
 deps = ["Pkg"]
 git-tree-sha1 = "acf12eccb390a78653ee805cd527898f01f78a85"
@@ -798,6 +803,7 @@ version = "0.6.1"
 deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+
 [[TikzGraphs]]
 deps = ["LaTeXStrings", "LightGraphs", "TikzPictures"]
 git-tree-sha1 = "48932ba660bc8cefc0aa9519ba79d63082aea892"
diff --git a/docs/src/lecture_08/ScalarReverseDiff.jl b/docs/src/lecture_08/ScalarReverseDiff.jl
@@ -57,6 +57,7 @@ cscheme = cgrad(:RdYlBu_5, rev=true)
 p1 = contour(-4:0.1:4, -2:0.1:2, g, fill=true, c=cscheme, xlabel="x", ylabel="y")
 display(p1)
 
+
 function descend(f::Function, λ::Real, args::Real...)
     Δargs = gradient(f, args...)
     args .- λ .* Δargs
diff --git a/docs/src/lecture_08/graphdiff.jl b/docs/src/lecture_08/graphdiff.jl
@@ -33,6 +33,7 @@ add_edge!(g, 2, 4)	# y ->  *
 add_edge!(g, 3, 5)	# sin(x) -> sin(x) + x*y
 add_edge!(g, 4, 5)  # x*y -> sin(x) + x*y
 add_edge!(g, 5, 6)  # sin(x) + x*y -> z
+
 #add_edge!(g, 5, 7,)	# ∂z/∂h₃ -> z
 #add_edge!(g, 4, 8,)	# ∂h₃/∂h₂ -> h₂ = *
 #add_edge!(g, 8, 9,)	# ∂z/∂h₂ -> ∂h₃/∂h₂
@@ -54,7 +55,6 @@ add_edge!(g, 5, 6)  # sin(x) + x*y -> z
 # 	markercolor = [vertices[i].color for i in 1:n],
 # )
 
-
 for n in [6, 7, 9, 11, 13, 14]
     #graphplot(adjacency_matrix(g)[1:n, 1:n],
 	#	names = [vertices[i].name for i in 1:n],
diff --git a/docs/src/lecture_08/graphdiff_11.svg b/docs/src/lecture_08/graphdiff_11.svg
diff --git a/docs/src/lecture_08/graphdiff_13.svg b/docs/src/lecture_08/graphdiff_13.svg
diff --git a/docs/src/lecture_08/graphdiff_14.svg b/docs/src/lecture_08/graphdiff_14.svg
@@ -367,4 +367,4 @@
 <path style="fill:none;stroke-width:0.79701;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;" d="M 39.313937 9.643031 L 4.642062 113.646938 " transform="matrix(1,0,0,-1,37.022,222.561)"/>
 <path style="fill:none;stroke-width:0.6376;stroke-linecap:round;stroke-linejoin:round;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;" d="M -1.552983 2.071716 C -1.424507 1.295934 -0.00141164 0.129794 0.387715 0.0000862163 C -0.00141082 -0.129624 -1.423263 -1.294537 -1.554205 -2.07032 " transform="matrix(-0.31621,-0.94864,-0.94864,0.31621,41.66565,108.91465)"/>
 </g>
-</svg>
+</svg>
diff --git a/docs/src/lecture_08/graphdiff_9.svg b/docs/src/lecture_08/graphdiff_9.svg
@@ -215,4 +215,4 @@
 <path style="fill:none;stroke-width:0.79701;stroke-linecap:butt;stroke-linejoin:miter;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;" d="M 39.313937 9.643031 L 4.642062 113.646938 " transform="matrix(1,0,0,-1,37.022,222.561)"/>
 <path style="fill:none;stroke-width:0.6376;stroke-linecap:round;stroke-linejoin:round;stroke:rgb(0%,0%,0%);stroke-opacity:1;stroke-miterlimit:10;" d="M -1.552983 2.071716 C -1.424507 1.295934 -0.00141164 0.129794 0.387715 0.0000862163 C -0.00141082 -0.129624 -1.423263 -1.294537 -1.554205 -2.07032 " transform="matrix(-0.31621,-0.94864,-0.94864,0.31621,41.66565,108.91465)"/>
 </g>
-</svg>
+</svg>
diff --git a/docs/src/lecture_08/jax-N=10000_M=1.txt b/docs/src/lecture_08/jax-N=10000_M=1.txt
@@ -0,0 +1,2 @@
+JAX (forward): 0.05906954288482666
+JAX (reverse): 0.004635481834411621
diff --git a/docs/src/lecture_08/jax-N=1000_M=1.txt b/docs/src/lecture_08/jax-N=1000_M=1.txt
@@ -0,0 +1,2 @@
+JAX (forward): 0.005011649131774903
+JAX (reverse): 0.003611617088317871
diff --git a/docs/src/lecture_08/jax-N=100_M=1.txt b/docs/src/lecture_08/jax-N=100_M=1.txt
@@ -0,0 +1,2 @@
+JAX (forward): 0.00287567138671875
+JAX (reverse): 0.0031633067131042482
diff --git a/docs/src/lecture_08/jax-N=100_M=100.txt b/docs/src/lecture_08/jax-N=100_M=100.txt
@@ -0,0 +1,2 @@
+JAX (forward): 0.002484147548675537
+JAX (reverse): 0.0030963587760925292
diff --git a/docs/src/lecture_08/jax-N=100_M=1000.txt b/docs/src/lecture_08/jax-N=100_M=1000.txt
@@ -0,0 +1,2 @@
+JAX (forward): 0.004091489315032959
+JAX (reverse): 0.007975153923034668
diff --git a/docs/src/lecture_08/julia-N=10000_M=1.txt b/docs/src/lecture_08/julia-N=10000_M=1.txt
@@ -0,0 +1,3 @@
+ForwardDiff.jl: 2.940003313e7
+ReverseDiff.jl: 271417.83
+Zygote.jl:  48022.63
diff --git a/docs/src/lecture_08/julia-N=1000_M=1.txt b/docs/src/lecture_08/julia-N=1000_M=1.txt
@@ -0,0 +1,3 @@
+ForwardDiff.jl: 381397.27
+ReverseDiff.jl: 68533.79
+Zygote.jl:  6498.564999999999
diff --git a/docs/src/lecture_08/julia-N=100_M=1.txt b/docs/src/lecture_08/julia-N=100_M=1.txt
@@ -0,0 +1,3 @@
+ForwardDiff.jl: 11961.82
+ReverseDiff.jl: 59821.06
+Zygote.jl:  1546.8603333333335
diff --git a/docs/src/lecture_08/julia-N=100_M=100.txt b/docs/src/lecture_08/julia-N=100_M=100.txt
@@ -0,0 +1,3 @@
+ForwardDiff.jl: 215920.0
+ReverseDiff.jl: 794220.7
+Zygote.jl:  20034.84
diff --git a/docs/src/lecture_08/julia-N=100_M=1000.txt b/docs/src/lecture_08/julia-N=100_M=1000.txt
@@ -0,0 +1,3 @@
+ForwardDiff.jl: 2.92624655e6
+ReverseDiff.jl: 1.179761551e7
+Zygote.jl:  173701.21
diff --git a/docs/src/lecture_08/lab.md b/docs/src/lecture_08/lab.md
@@ -16,6 +16,7 @@ to run `f` again, this time seeding the second argument with `Dual(y,one(y))`.
 Hence, we have to evaluate `f` *twice* if we want derivatives w.r.t to both its
 arguments which means that forward differentiation scales as $O(N)$ where $N$ is
 the number of inputs to `f`.
+
 ```julia
 dfdx = f(Dual(x,one(x)), Dual(y,zero(y)))
 dfdy = f(Dual(x,zero(x)), Dual(y,one(y)))
@@ -72,6 +73,7 @@ a = x*y               # da/dx = y;     da/dy = x
 b = sin(x)            # db/dx = cos(x)
 z = a + b             # dz/da = 1;     dz/db = 1
 ```
+
 ![graph](graph.png)
 
 In the graph you can see that the variable `x` can directly affect `b` and `a`.
@@ -171,6 +173,7 @@ y.children
 <header class="admonition-header">Exercise</header>
 <div class="admonition-body">
 ```
+
 Implement the two remaining rules for `+` and `sin` by overloading the
 appropriate methods like we did for `*`. First you have to compute the tracked
 forward pass, and then register the local derivatives in the children of your
@@ -239,6 +242,7 @@ z.grad = 1.0
 dx = accum!(x)
 dx ≈ y.data + cos(x.data)
 ```
+
 By accumulating the gradients for `x`, the gradients in the sub-tree connected
 to `x` will be evaluated. The parts of the tree that are only connected to `y`
 stay untouched.
@@ -309,6 +313,7 @@ where $\lambda$ is the learning rate that has to be tuned manually.
 <header class="admonition-header">Exercise</header>
 <div class="admonition-body">
 ```
+
 Implement a function `descend` performs one step of Gradient Descent (GD) on a
 function `f` with an arbitrary number of inputs. For GD you also have to
 specify the learning rate $\lambda$ so the function signature should look like
@@ -331,6 +336,7 @@ nothing # hide
 ```@raw html
 </p></details>
 ```
+
 Running one `descend` step should result in two new inputs with a smaller output
 for `g`
 ```@repl lab08
@@ -640,6 +646,7 @@ BenchmarkTools.Trial: 10000 samples with 6 evaluations.
 
  Memory estimate: 3.08 KiB, allocs estimate: 31.
 ```
+
 Even for this tiny example we are already 10 times faster than with the naively
 vectorized approach!
 
diff --git a/docs/src/lecture_08/lecture.md b/docs/src/lecture_08/lecture.md
@@ -4,6 +4,7 @@ using Plots
 # Automatic Differentiation
 
 ## Motivation
+
 - It supports a lot of modern machine learning by allowing quick differentiation of complex mathematical functions. The 1st order optimization methods are ubiquitous in finding parameters of functions (not only in  deep learning).
 - AD is interesting to study from the implementation perspective. There are different takes on it with different trade-offs and Julia offers many implementations (some of them are not maintained anymore).
 - We (authors of this course) believe that it is good to understand (at least roughly), how the methods work in order to use them effectively in our work.
@@ -13,6 +14,7 @@ using Plots
 
 The differentiation is routine process, as most of the time we break complicated functions down into small pieces that we know, how to differentiate and from that to assemble the gradient of the complex function back. Thus, the essential piece is the differentiation of the composed function ``f: \mathbb{R}^n \rightarrow \mathbb{R}^m``
 
+
 ```math
 f(x) = f_1(f_2(f_3(\ldots f_n(x)))) = (f_1 \circ f_2 \circ \ldots \circ f_n)(x)
 ``` 
@@ -91,6 +93,7 @@ Notice that
 The above is an idealized computation. The real implementation is a bit different, as we will see later.
 
 ### Implementation of the forward mode using Dual numbers
+
 Forward modes need to keep track of the output of the function and of the derivative at each computation step in the computation of the complicated function ``f``. This can be elegantly realized with a [**dual number**](https://en.wikipedia.org/wiki/Dual_number), which are conceptually similar to complex numbers, but instead of the imaginary number ``i`` dual numbers use ``\epsilon`` in its second component:
 ```math
 x = v + \dot v \epsilon,
@@ -279,6 +282,7 @@ We then add a node calculating ``\frac{\partial h_3}{\partial h_2}`` for which w
 
 We continue with the same process with ``\frac{\partial h_3}{\partial h_1}``, which we again combine with ``\frac{\partial z}{\partial h_1}`` to obtain ``\frac{\partial z}{\partial h_1}``. Continuing the reverse diff we obtain the final graph
 
+
 ![diff graph](graphdiff_14.svg) 
 
 containing the desired nodes ``\frac{\partial z}{\partial x}`` and ``\frac{\partial z}{\partial y}``. This computational graph can be passed to the compiler to compute desired values.
@@ -449,6 +453,7 @@ Disadvantages:
 !!! info
     The difference between tracking and graph-based AD systems is conceptually similar to interpreted and compiled programming languages. Tracking AD systems interpret the time while computing the gradient, while graph-based AD systems compile the computation of the gradient.
 
+
 ## ChainRules
 From our discussions about AD systems so far we see that while the basic, *engine*, part is relatively straightforward, it devil is in writing the rules prescribing the computation of gradients. These rules are needed for every system whether it is graph based, tracking, or Wengert list based. ForwardDiff also needs a rule system, but rules are a bit different (as they are pushing the gradient forward rather than pulling it back). It is obviously a waste of effort for each AD system to have its own set of rules. Therefore the community (initiated by Lyndon White backed by [Invenia](https://github.com/invenia)) have started to work on a unified system to express differentiation rules, such that they can be shared between systems. So far, they are supported by `Zygote.jl`, `Nabla.jl`, `ReverseDiff.jl` and `Diffractor.jl`, suggesting that the unification approach is working.
 
@@ -548,10 +553,83 @@ What if we somehow be able to talk to the compiler and get this form from him?
 
 * [simplest viable implementation](https://juliadiff.org/ChainRulesCore.jl/dev/autodiff/operator_overloading.html#ReverseDiffZero)
 
+## ChainRules
+From our discussions about AD systems so far we see that while the basic, *engine*, part is relatively straightforward, it devil is in writing the rules prescribing the computation of gradients. These rules are needed for every system whether it is graph based, tracking, or wengert list based. ForwardDiff also needs a rule system, but rules are a bit different (as they are pushing the gradient forward rather than pulling it back). It is obviously a waste of effort for each AD system to have its own set of rules. Hence the community (initiated by Lyndon White and hence probably funded mainly by Invenia) have started to work on a unified system to express differentiation rules, such that they can be shared between systems. So far, they are supported by `Zygote.jl`, `Nabla.jl`, `ReverseDiff.jl` and `Diffractor.jl`, hence the approach seems to work.
+
+The definition of reverse diff rules follows the idea we have nailed above (we refer readers interested in forward diff rules [to official documentation](https://juliadiff.org/ChainRulesCore.jl)).
+
+`ChainRules` defines the reverse rules for function `foo` in a function `rrule` with the following signature
+```julia
+function rrule(::typeof(foo), args...; kwargs...)
+    ...
+    return y, pullback
+end
+```
+where
+- the first argument `::typeof(foo)` allows to dispatch on the function for which the rules is written
+- the output of function `foo(args...)` is returned as the first argument
+- `pullback(Δy)` takes the gradient of upstream functions with respect to the output of `foo(args)` and returns it multiplied by the jacobian of the output of `foo(args)` with respect to parameters of the function itself (recall the function can have paramters, as it can be a closure or a functor), and with respect to the arguments.
+```julia
+function pullback(Δy)
+    ...
+    return ∂self, ∂args...
+end
+```
+Notice that key-word arguments are not differentiated. This is a design decision with the explanation that parametrize the function, but most of the time, they are not differentiable.
+
+`ChainRules.jl` provides support for lazy (delayed) computation using `Thunk`. Its argument is a function, which is not evaluated until `unthunk` is called. There is also a support to signal that gradient is zero using `ZeroTangent` (which can save valuable memory) or to signal that the gradient does not exist using `NoTangent`. 
+
+How can we use ChainRules to define rules for our AD system?
+Let's first observe the output
+```julia
+using ChainRulesCore, ChainRules
+r, g = rrule(*, rand(2,2), rand(2,2))
+g(r)
+```
+
+With that, we can extend our AD system as follows
+```julia
+import Base: *, +, -
+for f in [:*, :+, :-]
+    @eval function $(f)(A::TrackedMatrix, B::AbstractMatrix)
+       C, pullback = rrule($(f), value(A), B)
+       C = track(C)
+       push!(A.tape, (C, Δ -> pullback(Δ)[2]))
+       C
+    end
 
-### How to test
-A tricky example with odd function
+    @eval function $(f)(A::AbstractMatrix, B::TrackedMatrix)
+       C, pullback = rrule($(f), A, value(B))
+       C = track(C)
+       push!(B.tape, (C, Δ -> pullback(Δ)[3]))
+       C
+    end
 
+    @eval function $(f)(A::TrackedMatrix, B::TrackedMatrix)
+       C, pullback = rrule($(f), value(A), value(B))
+       C = track(C)
+       push!(A.tape, (C, Δ -> pullback(Δ)[2]))
+       push!(B.tape, (C, Δ -> pullback(Δ)[3]))
+       C
+    end
+end
+```
+and we need to modify our `accum!` code to `unthunk` if needed
+```julia
+function accum!(A::TrackedArray)
+    isempty(A.tape) && return(A.deriv)
+    A.deriv .= sum(unthunk(g(accum!(r))) for (r, g) in A.tape)
+end
+```
+
+```
+A = track(rand(4,4))
+B = rand(4,1)
+C = rand(4,1)
+R = A * B + C
+R.deriv .= 1
+accum!(A)
+```
 
 ### Sources for this lecture
 * Mike Innes' [diff-zoo](https://github.com/MikeInnes/diff-zoo)
diff --git a/docs/src/lecture_09/argcheck.jl b/docs/src/lecture_09/argcheck.jl
@@ -0,0 +1,91 @@
+using MacroTools
+using IRTools
+
+using IRTools: branches, block, empty!, evalir, func, branch!, block, IR, @dynamo, xcall
+
+function _mark(label, ex)
+    label isa Symbol || error("label has to be a Symbol")
+    return Expr(
+        :block,
+        Expr(:meta, :begin_optional, label),
+        esc(ex),
+        Expr(:meta, :end_optional, label),
+    )
+end
+
+macro mark(label, ex)
+    _mark(label, ex)
+end
+
+
+foo(x) = bar(baz(x))
+
+function bar(x)
+    @mark print iseven(x) && println("The input is even.")
+    x
+end
+
+function baz(x)
+    @mark print x<0 && println("The input is negative.")
+    x
+end
+
+
+
+isbegin(e::Expr) = Meta.isexpr(e,:meta) && e.args[1]===:begin_optional
+isend(e::Expr) = Meta.isexpr(e,:meta) && e.args[1]===:end_optional
+
+
+skip(f::Core.IntrinsicFunction, args...) = f(args...)
+skip(f::Core.Builtin, args...) = f(args...)
+
+@dynamo function skip(args...)
+    ir = IR(args...)
+    delete_line = false
+    local orig
+
+    for (x,st) in ir
+        is_begin = isbegin(st.expr)
+        is_end   = isend(st.expr)
+
+        if is_begin
+            delete_line = true
+        end
+    
+        if is_begin
+            orig = block(ir,x)
+        elseif is_end
+            dest = block(ir,x)
+            if orig != dest
+                empty!(branches(orig))
+                branch!(orig,dest)
+            end
+        end
+
+        if delete_line
+            delete!(ir,x)
+        end
+
+        if is_end
+            delete_line = false
+        end
+
+        if haskey(ir,x) && Meta.isexpr(st.expr,:call)
+            ir[x] = IRTools.xcall(skip, st.expr.args...)
+        end
+    end
+    return ir
+end
+
+function skip(ex::Expr)
+end
+
+macro skip(ex)
+    ex.head == :call || error("Input expression has to be a `:call`.")
+    return xcall(skip, ex.args...)
+end
+
+display(@code_ir foo(-2))
+display(@code_ir skip(foo,-2))
+display(foo(-2))
+@skip foo(-2)
diff --git a/docs/src/lecture_09/calls.jl b/docs/src/lecture_09/calls.jl
diff --git a/docs/src/lecture_09/codeinfo.jl b/docs/src/lecture_09/codeinfo.jl
diff --git a/docs/src/lecture_09/irtools.jl b/docs/src/lecture_09/irtools.jl
diff --git a/docs/src/lecture_09/irtools.md b/docs/src/lecture_09/irtools.md
diff --git a/docs/src/lecture_09/lab.md b/docs/src/lecture_09/lab.md
diff --git a/docs/src/lecture_09/lecture.md b/docs/src/lecture_09/lecture.md
diff --git a/docs/src/lecture_09/timer.jl b/docs/src/lecture_09/timer.jl

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+JAX (forward): 0.05906954288482666`
	`2`	`+JAX (reverse): 0.004635481834411621`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+JAX (forward): 0.005011649131774903`
	`2`	`+JAX (reverse): 0.003611617088317871`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+JAX (forward): 0.00287567138671875`
	`2`	`+JAX (reverse): 0.0031633067131042482`