WIP: Eager mode

src/TensorFlow.jl

@@ -141,23 +141,13 @@ function deallocator(data, len, arg)
 
 end
 
-struct Context
-    attrs::Dict
-end
-
-Context() = Context(Dict())
-
-struct ContextStack
-    contexts::Vector{Context}
-end
-
-ContextStack() = ContextStack(Context[])
-
-const global_context = ContextStack()
+include("context.jl")
 
 function __init__()
     c_deallocator[] = @cfunction(deallocator, Cvoid, (Ptr{Cvoid}, Csize_t, Ptr{Cvoid}))
-    push!(global_context, default_context())
+    for context in default_context()
+        push!(global_context, context)
+    end
 end
 
 function load_python_process(;force_reload=false)

src/context.jl

@@ -0,0 +1,44 @@
+abstract type Context
+end
+
+struct ContextStack
+    contexts::Vector{Context}
+end
+
+ContextStack() = ContextStack(Context[])
+
+function Base.push!(stack::ContextStack, context::Context)
+    push!(stack.contexts, context)
+end
+
+function Base.pop!(stack::ContextStack)
+    pop!(stack.contexts)
+end
+
+function default_context()
+    return [ExecutionMode(eager=false)]
+end
+
+function context_value(context_type)
+    return global_context[context_type]
+end
+
+function Base.getindex(c::ContextStack, context_type)
+    value = nothing
+    for context in c.contexts
+        if isa(context, context_type)
+            value = context
+        end
+    end
+    return value
+end
+
+function with_context(block, ctx)
+    push!(global_context, ctx)
+    res = block()
+    pop!(global_context)
+    return res
+end
+
+
+const global_context = ContextStack()

src/eager.jl

@@ -1,4 +1,4 @@
-mutable struct EagerContext
+mutable struct EagerContext <: Context
     ptr::Ptr{Cvoid}
 end
 
@@ -121,9 +121,10 @@ end
 
 function EagerOp(op_name)
     if get_eager_context() === nothing
-        ctx = Context()
-        ctx.attrs["eager_context"] = EagerContext()
-        push!(global_context, ctx)
+        # ctx = Context()
+        # ctx.attrs["eager_context"] = EagerContext()
+        # push!(global_context, ctx)
+        push!(global_context, EagerContext())
     end
     ctx = get_eager_context()
     status = Status()
@@ -245,7 +246,6 @@ function clear_caches(ctx::EagerContext)
     @tfcall(:TFE_ContextClearCaches, Cvoid, (Ptr{Cvoid},), ctx)
 end
 
-
 function num_dims(h::TensorHandle)
     status = Status()
     res = @tfcall(:TFE_TensorHandleNumDims, Cint, (Ptr{Cvoid}, Ptr{Cvoid}), h, status)
@@ -335,52 +335,24 @@ function inplace_sub(x, y)
     Ops.inplace_sub(x, i, y)
 end
 
-function Base.push!(stack::ContextStack, context::Context)
-    push!(stack.contexts, context)
-end
-
-function Base.pop!(stack::ContextStack)
-    pop!(stack.contexts)
+struct ExecutionMode <: Context
+    eager::Bool
 end
 
-function default_context()
-    context = Context()
-    context.attrs["eager"] = false
-    return context
-end
+ExecutionMode(;eager=true) = ExecutionMode(eager)
 
 function enable_eager_execution()
-    context = Context()
-    context.attrs["eager"] = true
-    push!(global_context, context)
+    # context = Context()
+    # context.attrs["eager"] = true
+    # push!(global_context, context)
+    push!(global_context, ExecutionMode(eager=true))
     return nothing
 end
 
-function Base.getindex(c::ContextStack, name)
-    value = nothing
-    for context in c.contexts
-        if name in keys(context.attrs)
-            value = context.attrs[name]
-        end
-    end
-    return value
-end
-
-function context_value(name)
-    return global_context[name]
-end
-
 function in_eager_mode()
-    return context_value("eager")::Bool
-end
-
-function with_context(block, ctx)
-    push!(global_context, ctx)
-    res = block()
-    pop!(global_context)
-    return res
+    return context_value(ExecutionMode).eager
 end
 
 function get_eager_context()
-    return context_value("eager_context")
+    return context_value(EagerContext)
 end

src/summary_writer.jl

@@ -4,7 +4,7 @@ import ..TensorFlow
 const tf = TensorFlow
 import ..TensorFlow: tensorflow, Graph, get_def_graph, @py_proc
 
-struct FileWriter
+struct FileWriter <: tf.Context
     file_handle
     logdir::String
 end
@@ -100,19 +100,15 @@ function Base.write(writer::FileWriter, graph::Graph)
 end
 
 function set_default(writer::FileWriter)
-    context = tf.Context()
-    context.attrs["default_file_writer"] = writer
-    push!(tf.global_context, context)
+    push!(tf.global_context, writer)
 end
 
 function with_default(writer::FileWriter, block)
-    context = tf.Context()
-    context.attrs["default_file_writer"] = writer
-    tf.with_context(block, context)
+    tf.with_context(block, writer)
 end
 
 function get_default_file_writer()
-    return tf.global_context["default_file_writer"]
+    return tf.context_value(FileWriter)
 end
 
 function record_summary(summary_pb; step=0)
@@ -121,7 +117,6 @@ function record_summary(summary_pb; step=0)
     write(writer, summary_pb, step)
 end
 
-
 function Base.close(writer::FileWriter)
     close(writer.file_handle)
     nothing

src/tape.jl

@@ -17,17 +17,26 @@ end
 
 Tape(;kwargs...) = Tape(Dict{TensorHandle, TapeNode}(), Dict(kwargs...))
 
+struct TapeContext <: Context
+    tape::Union{Tape, Nothing}
+end
+
 function set_tape(new_tape=nothing)
     if new_tape === nothing
         new_tape = Tape()
     end
-    context = Context()
-    context.attrs["tape"] = new_tape
-    push!(global_context, context)
+    push!(global_context, TapeContext(new_tape))
     return new_tape
 end
 
-get_tape() = global_context["tape"]
+function get_tape()
+    tape_context = context_value(TapeContext)
+    if tape_context === nothing
+        return nothing
+    else
+        return tape_context.tape
+    end
+end
 
 function add_node(t, node)
     tape = get_tape()
@@ -128,9 +137,7 @@ end)
 end)
 
 function with_no_grad(f)
-    context = Context()
-    context.attrs["tape"] = nothing
-    res = with_context(f, context)
+    res = with_context(f, TapeContext(nothing))
     return res
 end