WIP: Connecting to ChainRulesCore for JuliaLang AD compat

julia/Project.toml

@@ -3,14 +3,18 @@ uuid = "bb22f25d-cb49-471c-b017-930e329a2928"
 version = "0.1.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CombinedParsers = "5ae71ed2-6f8a-4ed1-b94f-e14e8158f19e"
 
 [compat]
+ChainRulesCore = "^1.0"
 CombinedParsers = "^0.2"
+Zygote = "^0.6.22"
 julia = "^1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Zygote"]

julia/README.md

@@ -2,7 +2,7 @@
 
 DexCall provides a mechanism for calling dex-lang code from JuliaLang.
 Three main mechanism are provided for this: `evaluate`, `DexModule` and the `dex_func` string macro.
-Two helper methods are also provided: `juliaize` and `NativeFunction`.
+Several helper methods are also provided: `juliaize`, `dexize`, and `NativeFunction`.
 
 ## `evaluate`: just run a single Dex expression.
 `evaluate` takes in a Dex expression as a string and runs it, returning a `Atom` (see below).
@@ -53,7 +53,7 @@ julia> m.addTwo(m.y)
 "[44., 44., 44.]"
 ```
 
-## Atoms: `juliaize`, `NativeFunction`
+## Atoms: `juliaize`, `dexize` and `NativeFunction`
 
 `evaluate` and the contents of a `DexModule` are returned as `Atom`s.
 These can be displayed, but not much else.
@@ -87,6 +87,8 @@ julia> typeof(convert(Int64, m.x))
 Int64
 ```
 
+The inverse of `juliaize` is `dexize`, it is currently very limited.
+
 To convert function `Atom`s into something you can execute as if it was a regular julia function use `NativeFunction`.
 This will compile it and handing inputs and outputs without needing to del with `Atom`s directly.
 

julia/src/DexCall.jl

@@ -1,14 +1,16 @@
 "Calling Dex from Julia"
 module DexCall
+    using ChainRulesCore
     using CombinedParsers
     using CombinedParsers.Regexp
 
-    export evaluate, DexError, DexModule, juliaize, NativeFunction, @dex_func_str
+    export evaluate, DexError, DexModule, dexize, juliaize, NativeFunction, @dex_func_str
 
     include("api_types.jl")
     include("api.jl")
     include("evaluate.jl")
     include("native_function.jl")
+    include("chainrules.jl")
 
     # use this to disable free'ing haskell objects after we have closed the RTS
     const NO_FREE = Ref(false)

julia/src/api_types.jl

@@ -55,11 +55,3 @@ function CAtom(atm::Ptr{HsAtom})
     iszero(success) && throw_from_dex()
     return result[]
 end
-
-"""
-    juliaize(x)
-
-Get the corresponding Julia object from some output of Dex.
-"""
-juliaize(x::CAtom) = bust_union(x)
-juliaize(x::Ptr{HsAtom}) = juliaize(CAtom(x))

julia/src/chainrules.jl

@@ -0,0 +1,51 @@
+
+function ChainRulesCore.frule((_, ẋ), f::Atom, x::Atom)
+    ẋ isa Atom || throw(DomainError(ẋ, "Tangent to an Atom must be an Atom"))
+    env = f.ctx
+    env = insert(env, "f", f.ptr)
+    env = insert(env, "dx", ẋ.ptr)
+    env = insert(env, "x", x.ptr)
+
+    m = DexModule(raw"""
+    (y, pushforward) = linearize f x
+    dy = pushforward dx
+    """,
+    env
+    )
+    return m.y, m.dy
+end
+
+function ChainRulesCore.rrule(f::Atom, x::Atom)
+    env = f.ctx
+    env = insert(env, "f", f.ptr)
+    env = insert(env, "x", x.ptr)
+
+    m = DexModule(raw"""
+    (y, pushforward) = linearize f x
+    pullback = transposeLinear pushforward
+    """,
+    env
+    )
+
+    # It is important that we close over `m` as otherwise the env may be destroyed by GC
+    pullback(x̄::Atom)= (NoTangent(), m.pullback(x̄))
+    return m.y, pullback
+end
+
+ChainRulesCore.frule((_, ẋ), ::typeof(juliaize), x) = juliaize(x), juliaize(ẋ)
+function ChainRulesCore.rrule(::typeof(juliaize), x::Atom)
+    env= x.ctx
+
+    # pullback must take a julia typed cotangent and give back a dex typed cotangent
+    juliaize_pullback(ȳ) = (NoTangent(), dexize(ȳ, env))
+    return juliaize(x), juliaize_pullback
+end
+
+
+ChainRulesCore.frule((_, ẋ), ::typeof(dexize), x) = dexize(x), dexize(ẋ)
+function ChainRulesCore.rrule(::typeof(dexize), x)
+    # pullback must take a dex typed cotangent and give back a julia typed cotangent
+    dexize_pullback(ȳ) = (NoTangent(), juliaize(ȳ))
+    return dexize(x), dexize_pullback
+end
+

julia/src/evaluate.jl

@@ -13,9 +13,37 @@ end
 
 Base.show(io::IO, atom::Atom) = show(io, print(atom.ptr))
 
+"""
+    juliaize(x)
+
+Get the corresponding Julia object from some output of Dex.
+"""
+juliaize(x::CAtom) = bust_union(x)
+juliaize(x::Ptr{HsAtom}) = juliaize(CAtom(x))
 juliaize(x::Atom) = juliaize(x.ptr)
 Base.convert(::Type{T}, atom::Atom) where {T<:Number} = convert(T, juliaize(atom))
 
+"""
+    dexize(x)
+
+Get the corresponding Dex object from some output of Julia.
+
+NB: this is currently a hack that goes via string processing.
+"""
+function dexize(x::Float32, _module=PRELUDE, env=_module)
+    isnan(x) && return evaluate("nan", _module, env)
+    x === Inf32 && return evaluate("infinity", _module, env)
+    x === -Inf32 && return evaluate("-infinity", _module, env)
+
+    str = repr(x)
+    if endswith(str, "f0")
+        evaluate(str[1:end-2], _module, env)
+    else
+        # convert "123f45" into "123 * (intpow 10.0 45)"
+        evaluate(replace(str, "f"=> " * (intpow 10.0 ") * ")", _module, env)
+    end
+end
+
 
 function (self::Atom)(args...)
     # TODO: Make those calls more hygenic
@@ -60,8 +88,8 @@ julia> m.y
 "84"
 ```
 """
-function DexModule(source::AbstractString)
-    ctx = dex_eval(PRELUDE, source)
+function DexModule(source::AbstractString, parent_ctx=PRELUDE)
+    ctx = dex_eval(parent_ctx, source)
     ctx == C_NULL && throw_from_dex()
     m =  DexModule(ctx)
     finalizer(m) do _m

julia/test/chainrules.jl

@@ -0,0 +1,31 @@
+const double_dex = evaluate(raw"\x:Float. 2.0 * x")
+
+@testset "frule: dexize, evaluate, juliaize" begin
+    a, ȧ = frule((NoTangent(), 10f0), dexize, 1.5f0)
+    b, ḃ = frule((NoTangent(), ȧ), double_dex, a)
+    c, ċ = frule((NoTangent(), ḃ), juliaize, b)
+    @test c === 3.0f0
+    @test ċ === 20f0
+end
+
+@testset "rrule: dexize, evaluate, juliaize" begin
+    x = 1.5f0
+    a, a_pb = rrule(dexize, x)
+    b, b_pb = rrule(double_dex, a)
+    c, c_pb = rrule(juliaize, b)
+
+    @test c === 3.0f0
+    c̄ = 10f0
+    _, b̄ = c_pb(c̄)
+    _, ā = b_pb(b̄)
+    _, x̄ = a_pb(ā)
+
+    @test x̄ === 20f0
+end
+
+@testset "Integration Test: Zygote.jl" begin
+    double_via_dex = juliaize ∘ double_dex ∘ dexize
+    y, pb= Zygote.pullback(double_via_dex, 1.5f0)
+    @test y == 3f0
+    @test pb(1f0) == (2f0,)
+end

julia/test/evaluate.jl

@@ -21,6 +21,16 @@
         @test juliaize(evaluate("IToW8 65")) === Int8(65)
     end
 
+    @testset "dexize" begin
+        @test juliaize(dexize(0f0)) === 0f0
+        @test juliaize(dexize(42f0)) === 42f0
+        @test juliaize(dexize(123f15)) === 123f15
+        @test dexize(123f15) isa DexCall.Atom
+        @test isnan(juliaize(dexize(NaN32)))
+        @test (juliaize(dexize(Inf32))) == Inf32
+        @test (juliaize(dexize(-Inf32))) == -Inf32
+    end
+
     @testset "Atom function call" begin
         m = DexModule("""
         def addOne (x: Float) : Float = x + 1.0

julia/test/runtests.jl

@@ -1,8 +1,11 @@
 using Test
 using DexCall
+using ChainRulesCore
+using Zygote  # for integration tests 
 
 @testset "DexCall" begin
     include("api.jl")
     include("evaluate.jl")
     include("native_function.jl")
+    include("chainrules.jl")
 end