Make thunked unthunk before broadcasting (#97)

* Make thunked unthunk (not extern) before broadcasting

* bump version

* Fix spacing

Co-Authored-By: Nick Robinson <npr251@gmail.com>

Co-authored-by: Nick Robinson <npr251@gmail.com>

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.5.3"
+version = "0.5.4"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

src/differentials/thunks.jl

@@ -1,17 +1,17 @@
 
 abstract type AbstractThunk <: AbstractDifferential end
 
-Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(extern(x))
+Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(unthunk(x))
 
 @inline function Base.iterate(x::AbstractThunk)
-    externed = extern(x)
-    element, state = iterate(externed)
-    return element, (externed, state)
+    val = unthunk(x)
+    element, state = iterate(val)
+    return element, (val, state)
 end
 
-@inline function Base.iterate(::AbstractThunk, (externed, state))
-    element, new_state = iterate(externed, state)
-    return element, (externed, new_state)
+@inline function Base.iterate(::AbstractThunk, (val, state))
+    element, new_state = iterate(val, state)
+    return element, (val, new_state)
 end
 
 #####
@@ -62,7 +62,7 @@ Do not use `@thunk` if this would be equal or more work than actually evaluating
 - The expression is merely wrapping something in a `struct`, such as `Adjoint(x)` or `Diagonal(x)`
 - The expression being itself a `thunk`
 - The expression being from another `rrule` or `frule` (it would be `@thunk`ed if required by the defining rule already)
-- There is only one derivative being returned, so from the fact that the user called `frule`/`rrule` 
+- There is only one derivative being returned, so from the fact that the user called `frule`/`rrule`
 they clearly will want to use that one.
 """
 struct Thunk{F} <: AbstractThunk

test/differentials/thunks.jl

@@ -35,4 +35,40 @@
             @test stackframe.file == Symbol(@__FILE__)
         end
     end
+
+
+    @testset "Broadcast" begin
+        @testset "Array" begin
+            was_unthunked = 0
+            array_thunk = @thunk begin
+                was_unthunked += 1;
+                [1.0, 2.0, 3.0]
+            end
+
+            was_unthunked = 0
+            @test array_thunk .+ fill(10, 3) .+  fill(10, 3) == [21.0, 22.0, 23.0]
+            @test was_unthunked == 1
+
+            was_unthunked = 0
+            @test array_thunk .+ 10.0 .+ 10.0 == [21.0, 22.0, 23.0]
+            @test was_unthunked == 1
+
+        end
+
+        @testset "Scalar" begin
+            was_unthunked=0
+            scalar_thunk = @thunk begin
+                was_unthunked += 1;
+                sqrt(4.0)
+            end
+
+            was_unthunked = 0
+            @test scalar_thunk .+ fill(10, 3) .+  fill(10, 3) == [22.0, 22.0, 22.0]
+            @test was_unthunked == 1
+
+            was_unthunked = 0
+            @test scalar_thunk .+ 10.0 .+ 10.0 == 22.0
+            @test was_unthunked == 1
+        end
+    end
 end