Make `extern` recursive?

[oxinabox]

it is not too hard to endup in a situtation where you have say
Thunk(()->Thunk(()->3)) ie., @thunk(@thunk(3))

One example (written in #30 style) is if you have

rrule(::typeof(f), x) = f(x), y->(NO_FIELDS, @thunk(3))
function rrule(::typeof(g), x)
    _, inner_pullback = rrule(f, x)
   function g_pullback(y)
        (NO_FIELDS, @thunk(g_pullback(y)[2])
   end
   return g(x), g_pullback
end

Especially with #30 where we (by nescity) throw a lot more @thunks around.

This is particularly annoying in test right now, and so probably in real use too.

---

I see two ways to resolve this:

1. make extern(::Thunk) recursive

this is just changing

@inline extern(x::Thunk) = x.f()

to

@inline extern(x::Thunk) = extern(x.f())

This doesn't help for (Casted(Zero()) or Casted(@thunk(3)) etc
but Casted could have the same change I guess,
or it could just or away #10

2. make all externing recursive

We would change current definitions of extern into
extern1
then change extern to call extern1 until the result is no longer a AbstractDifferential.
(either via recursion, or via a loop)

3. Make `Thunk(()->Thunk

Way harder than it looks.
Similarly does not apply to things other than thunks

---

Not sure on the implications of these on inlining,
I know inlining hates recursion.

Here is some test code I prepared earlier.

    @testset "Thunk" begin
        @test extern(Thunk(()->3)) == 3 == extern(@thunk(3))
        @test extern(Thunk(()->Thunk(()->3))) == 3 == extern(@thunk(@thunk(3)))
        @test @thunk(3) isa Thunk
    end

I think the first one is a good place to start.
And will likely PR that and rebase it into #30

[oxinabox]

Some of the tests right now rely on
extern(@thunk Wirtinger(...)) isa Wirtinger
but those could be change to be
just check the field of the thunk.
Or we could make thunks callable and calling the Thunk could call the inside.
So it wouild just extern once.
I think I would like thunks to be callable anyway.

[oxinabox]

I tried out option 1.

Seems to inline just fine

julia> @code_typed extern(@thunk(3))
CodeInfo(
1 ─     return 3
) => Int64

julia> @code_typed extern(@thunk(@thunk(3)))
CodeInfo(
1 ─     return 3
) => Int64

julia> @code_typed extern(@thunk(@thunk(Zero())))
CodeInfo(
1 ─     return false
) => Bool


julia> @code_typed extern(@thunk(@thunk(identity(3))
CodeInfo(
1 ─     return 3
) => Int64

julia> @code_typed extern(@thunk(@thunk(sum(1:3))))
CodeInfo(
1 ─     return 6
) => Int64

julia> @code_typed extern(@thunk(@thunk(rand())))
CodeInfo(
1 ── %1  = Random.GLOBAL_RNG::MersenneTwister
│    %2  = (Base.getfield)(%1, :idxF)::Int64
│    %3  = Random.MT_CACHE_F::Int64
...

[nickrobinson251]

Is inlining in anyway version dependent? E.g. does it ever work on Julia 1.3 but not 1.0?

[oxinabox]

That above was on 1.1. I could check on the LTS

Update: yes, it is also all good on the LTS

[simeonschaub]

Two additional thoughts:
Should extern also broadcast over arrays containing AbstractDifferentials?
wirtinger_primal and wirtinger_conjugate should probably also work in a similar matter. I recently ran into that with a Thunk returning Wirtinger, but should this stop being recursive once a Wirtinger object is encountered or should it also call extern on the components if those are differentials?

[oxinabox]

Should extern also broadcast over arrays containing AbstractDifferentials?

We should probably try and avoid creating those.
Not sure we can entirely. Though.

I recently ran into that with a Thunk returning Wirtinger, but should this stop being recursive once a Wirtinger object is encountered... ?

Yes @scalar_rule makes these I think?

I think extern(@thunk Wirtinger (...)) should be the same as extern(Wirtinger(...)) which is an error.
And I think that is right. Eyern gives back results in external types.

We might want something else like unthunk or resolve or something that is kinda similar but ends up with simpler in-domain types.
Less sure.