Julia operators are functions

As already implied by the syntax tree form of \(g(x)\), almost⁴ all elementary operators like *, +, etc are actually Julia functions with special infix syntax. They are defined as such and can be invoked with the “normal” function call syntax:

julia> x = 2;

julia> x + x
4

julia> +(x, x) # call function named "+" with parameter tuple "(x, x)"
4

(In fact, most of Julia is defined via its own syntax, something the language designers and user community are justifiably proud of.)

Julia functions are “overloadable” (or is it “overridable”? 😄)

Multiplying two integers is not the same as “multiplying” two strings⁵:

julia> x * x
4

julia> "x" * "x"
"xx"

and might be different from “multiplying” objects of a user-defined type. As I started to explain in my previous typing tutorial, this is made possible by using different parameter types for different function versions (“methods”, in Julia parlance). Internally, Julia keeps a list of all methods it is aware of for each operator/function and they can be examined using methods():

julia> methods(*)
# 358 methods for generic function "*":
[1] *(x::Bool, z::Complex{Bool}) in Base at complex.jl:282
[2] *(x::Bool, y::Bool) in Base at bool.jl:98
[3] *(x::Bool, y::T) where T<:AbstractFloat in Base at bool.jl:110
[4] *(x::Bool, z::Complex) in Base at complex.jl:289
[5] *(x::Bool, y::AbstractIrrational) in Base at irrationals.jl:139
[6] *(a::Float16, b::Float16) in Base at float.jl:392
[7] *(x::Float32, y::Float32) in Base at float.jl:398
[8] *(x::Float64, y::Float64) in Base at float.jl:399
[9] *(z::Complex{Bool}, x::Bool) in Base at complex.jl:283
... lots more ...

It is possible for a user to add more methods to such lists. Each row has a unique signature. Imagine defining your own “numbers”:

julia> struct MyNum <: Number
           v   ::Float64
       end

julia> MyNum(2.0) * MyNum(2.0) # this won't work
ERROR: MethodError: no method matching *(::MyNum, ::MyNum)
Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...) at operators.jl:529

julia> Base.:(*)(lhs ::MyNum, rhs ::MyNum) = MyNum(lhs.v * rhs.v) # define 'MyNum * MyNum', so now it will

julia> MyNum(2.0) * MyNum(2.0)
MyNum(4.0)

julia> methods(MyNum)
# 2 methods for generic function "(::Type)":
[1] MyNum(v::Float64) in Main at none:1
[2] MyNum(v) in Main at none:1

julia> methods(*)
# 359 methods for generic function "*": # <- note that count is up by one
[1] *(x::Bool, z::Complex{Bool}) in Base at complex.jl:282
...
[31] *(lhs::MyNum, rhs::MyNum) in Main at none:1
...

Note again how the new method doesn’t “belong” to MyNum but rather to a global method table associated with *.

Continuing, I would need to define a minimal collection of arithmetic operations in order for `MyNum’s to be useful as something resembling numbers:

# binary ops:

Base.:(+)(lhs ::MyNum, rhs ::MyNum) = MyNum(lhs.v + rhs.v)
Base.:(-)(lhs ::MyNum, rhs ::MyNum) = MyNum(lhs.v - rhs.v)

Base.:(*)(lhs ::MyNum, rhs ::MyNum) = MyNum(lhs.v * rhs.v)
Base.:(/)(lhs ::MyNum, rhs ::MyNum) = MyNum(lhs.v / rhs.v)

# unary ops:

Base.:(+)(x ::MyNum) = x
Base.:(-)(x ::MyNum) = MyNum(- x.v)

julia> MyNum(3.0) + MyNum(4.0)
MyNum(7.0)

So far, so good. Note that \(g(x)\) is generic enough so it “just works” for both plain floats and MyNums:

julia> g(2.0)
8.0

julia> g(MyNum(2.0))
MyNum(8.0)

Take note, this will happen again later. Of course, \(g(x)\) is downright trivial: it uses only * and + – this was truly a baby step.

Something is missing

My custom, soon-to-win-many-awards, number type is missing some boilerplate stuff. Trying another simple function fails; it flat out refuses to work with MyNums:

julia> f(x) = 1.0 / x
f (generic function with 1 method)

julia> f(MyNum(2.0))
ERROR: promotion of types Float64 and MyNum failed to change any arguments
Stacktrace:
 [1] sametype_error(::Tuple{Float64,MyNum}) at ./promotion.jl:308
 [2] not_sametype(::Tuple{Float64,MyNum}, ::Tuple{Float64,MyNum}) at ./promotion.jl:302
 [3] promote at ./promotion.jl:285 [inlined]
 [4] /(::Float64, ::MyNum) at ./promotion.jl:316
 [5] f(::MyNum) at ./none:1

I can see from the stacktrace that \(f(x)\) was indeed found and started, but something went wrong later: there was no suitable defition of / for parameter type pair (Float64, MyNum). Julia does not know to convert the 1.0 in 1.0 / x to a MyNum – the language does not have what’s called automatic promotion. As a result, MyNums don’t interoperate with the “stock” numbers particularly well.

julia> Base.:(/)(lhs ::Float64, rhs ::MyNum) = MyNum(lhs / rhs.v) # float / MyNum

julia> f(MyNum(2.0)) # f() works now
MyNum(0.5)

julia> h(x) = x / 3.0
h (generic function with 1 method)

julia> h(MyNum(2.0))
ERROR: promotion of types MyNum and Float64 failed to change any arguments
...

julia> Base.:(/)(lhs ::MyNum, rhs ::Float64) = MyNum(lhs.v / rhs) # MyNum / float

julia> h(MyNum(2.0)) # h() works now, too
MyNum(0.6666666666666666)

Base.:(/)(lhs ::MyNum,   rhs ::MyNum)   = MyNum(lhs.v / rhs.v)
Base.:(/)(lhs ::Float64, rhs ::MyNum)   = MyNum(lhs / rhs.v)
Base.:(/)(lhs ::MyNum,   rhs ::Float64) = MyNum(lhs.v / rhs)

Base.promote_rule(::Type{MyNum}, ::Type{<: Number}) = MyNum

Core AD idea #1: composition of steps

If you’ve made it through the preceeding sections, the following should now be seared into your subconscious: a function computes its output by forward-propagating inputs and intermediate expression values through its calculation graph. These elementary steps are +, *, cos, sin, etc.

Futhermore, these steps are not very “intrinsic” or somehow untouchable: I’ve shown that I could wrap values into a custom type and overload all the operations I am interested in. This works beacuse of Julia’s core paradigm of (multiple) dispatch on argument types.

Core AD idea #2: step instrumentation

Just one more step remaining. Why did I bother replicating standard arithmetic with MyNum? Sure, it was to understand better how some aspects of Julia worked. But note this: computing the derivative of a function is also a composable calculation and it can be done in parallel with computing the function value.

So… if I extend MyNum with another field to hold the current expression’s derivative, I can forward-propagate it in the same pass as the main" value:

struct Context <: Number
    v   ::Number
    ∂   ::Number
end

# binary ops:

Base.:(+)(lhs ::Context, rhs ::Context) = Context(lhs.v + rhs.v, lhs.∂ + rhs.∂)
Base.:(-)(lhs ::Context, rhs ::Context) = Context(lhs.v - rhs.v, lhs.∂ - rhs.∂)

Base.:(*)(lhs ::Context, rhs ::Context) = Context(lhs.v * rhs.v, lhs.v * rhs.∂ + lhs.∂ * rhs.v)
Base.:(/)(lhs ::Context, rhs ::Context) = Context(lhs.v / rhs.v, (lhs.∂ * rhs.v - lhs.v * rhs.∂) / rhs.v^2)

# unary ops:

Base.:(+)(x ::Context) = x
Base.:(-)(x ::Context) = Context(- x.v, - x.∂)

Here Context is the grown-up version of MyNum. It has been extended with field \(\partial\). The arithmetic overloads for Context continue to propagate expression values in \(v\) while some new logic propagates derivative values in \(\partial\).

Individual Context \(v\)-rules are composed to form a larger calculation such as myerf(x). Consistently with such pattern of composition, each \(\partial\)-rule can be seen to be an application of the chain rule for obtaining derivatives of the composition of functions. For example, since \[ \frac{\partial}{\partial x} (lhs + rhs)(x) = \frac{\partial}{\partial x} lhs(x) + \frac{\partial}{\partial x} rhs(x) \] the Context rule for + must be \[ +\big( \langle lhs.v, lhs.\partial \rangle, \langle rhs.v, rhs.\partial \rangle \big) = \big\langle lhs.v + rhs.v, lhs.\partial + rhs.\partial \big\rangle. \] The rule for sin, should it be needed, would be \[ \sin \big(\langle x.v, x.\partial \rangle \big) = \big\langle \sin(x.v), \cos(x.v) \partial \big\rangle, \] and so on.

What Context value should be fed as the initial input? Its \(v\)-field is x. The derivative of that with respect to x is 1.0, so to establish derivative propagation correctly it needs to be seeded with Context(x, 1.0).

And that’s pretty much it. Does it work? Let’s give it a try:

function derivative(f ::Function)
    return x ::Number -> f(Context(x, 1.0)).∂ # discard value, return derivative
end

∂ = derivative # add a nice-looking alias

And the moment of truth:

julia> include("multiple_dispatch/FAD.jl")
Main.FAD

julia> using Main.FAD # bring ∂, etc into this REPL

julia> x = 0.7;

julia> ∂(myerf)(x), 2/√π * exp(-x^2)
(0.6912748604105389, 0.6912748604105386)

Success!

As a different test, the chain rule \(\frac{\partial}{\partial x}(f \circ g)(x) = \frac{\partial}{\partial g} f(g(x)) \frac{\partial}{\partial x} g(x)\) should also hold for arbitrary “differentiable” Julia functions:

julia> f(x) = 1.0 / x
f (generic function with 1 method)

julia> g(x) = x * (x + x)
g (generic function with 1 method)

julia> x = 1.2345;

julia> ∂(f ∘ g)(x), ∂(f)(g(x)) * ∂(g)(x)
(-0.5315286974115385, -0.5315286974115385)

I think this is downright beautiful.

And last but not least, let me try an algorithm that can benefit from having both \(f(x)\) and \(\frac{\partial}{\partial x}f(x)\) available at the same time, e.g. Newton-Raphons root finder. Here is a prototype version performing \(x_{n+1} = x_n - \frac{f(x_n)}{\frac{\partial}{\partial x}f(x_n)}\) iterations until convergence within given tolerance:

function root_solve(f ::Function, x₀ ::Number; ϵ = 1e-8)
    i = 1
    while true
        ctx = f(Context(x₀, 1.0))
        println("[$i]: f($x₀)\t= $(ctx.v)")
        abs(ctx.v) < ϵ && break
        x₀ -= ctx.v / ctx.∂ # use both value and derivative
        i += 1
    end
    return x₀
end

If I extend my Context “algebra” with some more elementary operations

Base.sin(d ::Context) = Context(sin(d.v),   cos(d.v) * d.∂)
Base.cos(d ::Context) = Context(cos(d.v), - sin(d.v) * d.∂)

then I can try solving, say, trigonometric equations:

julia> root_solve(x -> sin(x) - cos(x), 0.0) # find a root of sin(x) = cos(x) starting from x₀ = 0
[1]: f(0.0) = -1.0
[2]: f(1.0) = 0.30116867893975674
[3]: f(0.782041901539138)   = -0.004746462127804163
[4]: f(0.7853981759997019)  = 1.7822277875723103e-8
[5]: f(0.7853981633974483)  = -1.1102230246251565e-16
0.7853981633974483

julia> π/4
0.7853981633974483

Nice. I find this all very satisfying. Especially given how few lines of Julia actually went into the final implementation.

For an industrial-strength implementation of this approach to AD, check out ForwardDiff.jl.