add Unthunk

Author: oxinabox

src/ChainRulesCore.jl

@@ -4,6 +4,7 @@ using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broad
 export frule, rrule
 export wirtinger_conjugate, wirtinger_primal, refine_differential
 export @scalar_rule, @thunk
+export unthunk,
 export extern, cast, store!
 export Wirtinger, Zero, DoesNotExist, Thunk, InplaceableThunk
 export NO_FIELDS

src/differential_arithmetic.jl

@@ -67,12 +67,12 @@ for T in (:AbstractThunk, :Any)
 end
 
 
-Base.:+(a::AbstractThunk, b::AbstractThunk) = extern(a) + extern(b)
-Base.:*(a::AbstractThunk, b::AbstractThunk) = extern(a) * extern(b)
+Base.:+(a::AbstractThunk, b::AbstractThunk) = unthunk(a) + unthunk(b)
+Base.:*(a::AbstractThunk, b::AbstractThunk) = unthunk(a) * unthunk(b)
 for T in (:Any,)
-    @eval Base.:+(a::AbstractThunk, b::$T) = extern(a) + b
-    @eval Base.:+(a::$T, b::AbstractThunk) = a + extern(b)
+    @eval Base.:+(a::AbstractThunk, b::$T) = unthunk(a) + b
+    @eval Base.:+(a::$T, b::AbstractThunk) = a + unthunk(b)
 
-    @eval Base.:*(a::AbstractThunk, b::$T) = extern(a) * b
-    @eval Base.:*(a::$T, b::AbstractThunk) = a * extern(b)
+    @eval Base.:*(a::AbstractThunk, b::$T) = unthunk(a) * b
+    @eval Base.:*(a::$T, b::AbstractThunk) = a * unthunk(b)
 end

src/differentials.jl

@@ -149,6 +149,18 @@ end
     return element, (externed, new_state)
 end
 
+"""
+    unthunk(x)
+
+On `AbstractThunk`s this removes 1 layer of thunking.
+On any other type, it is the identity operation.
+
+In contrast to `extern` this is nonrecursive.
+"""
+@inline unthunk(x) = x
+
+@inline extern(x::AbstractThunk) = extern(unthunk(x))
+
 #####
 ##### `Thunk`
 #####
@@ -213,9 +225,9 @@ end
 # have to define this here after `@thunk` and `Thunk` is defined
 Base.conj(x::AbstractThunk) = @thunk(conj(extern(x)))
 
-
 (x::Thunk)() = x.f()
-@inline extern(x::Thunk) = extern(x())
+@inline unthunk(x::Thunk) = x()
+
 
 Base.show(io::IO, x::Thunk) = println(io, "Thunk($(repr(x.f)))")
 
@@ -237,8 +249,8 @@ struct InplaceableThunk{T<:Thunk, F} <: AbstractThunk
     add!::F
 end
 
-(x::InplaceableThunk)() = x.val()
-@inline extern(x::InplaceableThunk) = extern(x.val)
+@inline unthunk(x::InplaceableThunk) = unthunk(x.val)
+(x::InplaceableThunk)() = unthunk(x)
 
 function Base.show(io::IO, x::InplaceableThunk)
     println(io, "InplaceableThunk($(repr(x.val)), $(repr(x.add!)))")

test/differentials.jl

@@ -45,6 +45,11 @@
             @test extern(@thunk(@thunk(3))) == 3
         end
 
+        @testset "unthunk" begin
+            @test unthunk(@thunk(3)) == 3
+            @test unthunk(@thunk(@thunk(3))) isa Thunk
+        end
+
         @testset "calling thunks should call inner function" begin
             @test (@thunk(3))() == 3
             @test (@thunk(@thunk(3)))() isa Thunk

test/rules.jl

@@ -134,11 +134,11 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
             fx, f_pushforward = res
             df(Δx, Δp) = f_pushforward(NamedTuple(), Δx, Δp)
 
-            df_dx::Thunk = df(one(x), Zero())
-            df_dp::Thunk = df(Zero(), one(p))
+            df_dx = df(one(x), Zero())
+            df_dp = df(Zero(), one(p))
             @test fx == f(x, p)  # Check we still get the normal value, right
-            @test df_dx() isa expected_type_df_dx
-            @test df_dp() isa expected_type_df_dp
+            @test df_dx isa expected_type_df_dx
+            @test df_dp isa expected_type_df_dp
 
 
             res = rrule(f, x, p)
@@ -147,8 +147,8 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
             dself, df_dx, df_dp = f_pullback(One())
             @test fx == f(x, p)  # Check we still get the normal value, right
             @test dself == NO_FIELDS
-            @test df_dx() isa expected_type_df_dx
-            @test df_dp() isa expected_type_df_dp
+            @test df_dx isa expected_type_df_dx
+            @test df_dp isa expected_type_df_dp
         end
     end
 end