checkpoint on getting tests to pass again

Author: marcoct

src/Gen.jl

@@ -64,9 +64,6 @@ include("dynamic/dynamic.jl")
 # static IR generative function
 include("static_ir/static_ir.jl")
 
-# optimization for built-in generative functions (dynamic and static IR)
-include("builtin_optimization.jl")
-
 # DSLs for defining dynamic embedded and static IR generative functions
 # 'Dynamic DSL' and 'Static DSL'
 include("dsl/dsl.jl")

src/builtin_optimization.jl

@@ -1,235 +0,0 @@
-#############################
-
-# primitives for in-place gradient accumulation
-
-function in_place_add!(param::Array, increment::Array, scale_factor::Real)
-    # NOTE: it ignores the scale_factor, because it is not a parameter...
-    # scale factors only affect parameters
-    # TODO this is potentially very confusing!
-    @simd for i in 1:length(param)
-        param[i] += increment[i]
-    end
-    return param
-end
-
-function in_place_add!(param::Array, increment::Array)
-    @inbounds @simd for i in 1:length(param)
-        param[i] += increment[i]
-    end
-    return param
-end
-
-function in_place_add!(param::Real, increment::Real, scale_factor::Real)
-    return param + increment
-end
-
-function in_place_add!(param::Real, increment::Real)
-    return param + increment
-end
-
-mutable struct ThreadsafeAccumulator{T}
-    value::T
-    lock::ReentrantLock
-end
-
-ThreadsafeAccumulator(value) = ThreadsafeAccumulator(value, ReentrantLock())
-
-# TODO not threadsafe
-function get_current_value(accum::ThreadsafeAccumulator)
-    return accum.value
-end
-
-function in_place_add!(param::ThreadsafeAccumulator{Real}, increment::Real, scale_factor::Real)
-    lock(param.lock)
-    try
-        param.value = param.value + increment * scale_factor
-    finally
-        unlock(param.lock)
-    end
-    return param
-end
-
-function in_place_add!(param::ThreadsafeAccumulator{Real}, increment::Real)
-    lock(param.lock)
-    try
-        param.value = param.value + increment
-    finally
-        unlock(param.lock)
-    end
-    return param
-end
-
-function in_place_add!(param::ThreadsafeAccumulator{<:Array}, increment, scale_factor::Real)
-    lock(param.lock)
-    try
-        @simd for i in 1:length(param.value)
-            param.value[i] += increment[i] * scale_factor
-        end
-    finally
-        unlock(param.lock)
-    end
-    return param
-end
-
-function in_place_add!(param::ThreadsafeAccumulator{<:Array}, increment)
-    lock(param.lock)
-    try
-        @simd for i in 1:length(param.value)
-            param.value[i] += increment[i]
-        end
-    finally
-        unlock(param.lock)
-    end
-    return param
-end
-
-#############################
-
-
-
-"""
-    set_param!(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, value)
-
-Set the value of a trainable parameter of the generative function.
-
-NOTE: Does not update the gradient accumulator value.
-"""
-function set_param!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, value)
-    return gf.params[name] = value
-end
-
-"""
-    value = get_param(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-
-Get the current value of a trainable parameter of the generative function.
-"""
-function get_param(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-    return gf.params[name]
-end
-
-"""
-    value = get_param_grad(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-
-Get the current value of the gradient accumulator for a trainable parameter of the generative function.
-
-Not threadsafe.
-"""
-function get_param_grad(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-    try
-	val = gf.params_grad[name] # the accumulator
-        return get_current_value(val)
-    catch KeyError
-        error("parameter $name not found")
-    end
-    return val
-end
-
-"""
-    zero_param_grad!(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-
-Reset the gradient accumlator for a trainable parameter of the generative function to all zeros.
-
-Not threadsafe.
-"""
-function zero_param_grad!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol)
-    gf.params_grad[name] = ThreadsafeAccumulator(zero(gf.params[name])) # TODO avoid allocation?
-    return gf.params_grad[name]
-end
-
-"""
-    set_param_grad!(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, grad_value)
-
-Set the gradient accumlator for a trainable parameter of the generative function.
-
-Not threadsafe.
-"""
-function set_param_grad!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, grad_value)
-    gf.params_grad[name] = ThreadsafeAccumulator(grad_value)
-    return grad_value
-end
-
-# TODO document me; it is threadsafe..
-function increment_param_grad!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, increment, scale_factor)
-    in_place_add!(gf.params_grad[name], increment, scale_factor)
-end
-
-# TODO document me; it is threadsafe..
-function increment_param_grad!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, increment)
-    in_place_add!(gf.params_grad[name], increment)
-end
-
-
-
-"""
-    init_param!(gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, value)
-
-Initialize the the value of a named trainable parameter of a generative function.
-
-Also generates the gradient accumulator for that parameter to `zero(value)`.
-
-Example:
-```julia
-init_param!(foo, :theta, 0.6)
-```
-"""
-function init_param!(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, name::Symbol, value)
-    set_param!(gf, name, value)
-    zero_param_grad!(gf, name)
-end
-
-get_params(gf::Union{DynamicDSLFunction,StaticIRGenerativeFunction}) = keys(gf.params)
-
-export set_param!, get_param, get_param_grad, zero_param_grad!, set_param_grad!, init_param!, increment_param_grad!
-
-#########################################
-# gradient descent with fixed step size #
-#########################################
-
-mutable struct FixedStepGradientDescentBuiltinDSLState
-    step_size::Float64
-    gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}
-    param_list::Vector
-end
-
-function init_update_state(conf::FixedStepGradientDescent,
-        gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, param_list::Vector)
-    FixedStepGradientDescentBuiltinDSLState(conf.step_size, gen_fn, param_list)
-end
-
-function apply_update!(state::FixedStepGradientDescentBuiltinDSLState)
-    for param_name in state.param_list
-        value = get_param(state.gen_fn, param_name)
-        grad = get_param_grad(state.gen_fn, param_name)
-        set_param!(state.gen_fn, param_name, value + grad * state.step_size)
-        zero_param_grad!(state.gen_fn, param_name)
-    end
-end
-
-####################
-# gradient descent #
-####################
-
-mutable struct GradientDescentBuiltinDSLState
-    step_size_init::Float64
-    step_size_beta::Float64
-    gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}
-    param_list::Vector
-    t::Int
-end
-
-function init_update_state(conf::GradientDescent,
-        gen_fn::Union{DynamicDSLFunction,StaticIRGenerativeFunction}, param_list::Vector)
-    GradientDescentBuiltinDSLState(conf.step_size_init, conf.step_size_beta,
-        gen_fn, param_list, 1)
-end
-
-function apply_update!(state::GradientDescentBuiltinDSLState)
-    step_size = state.step_size_init * (state.step_size_beta + 1) / (state.step_size_beta + state.t)
-    for param_name in state.param_list
-        value = get_param(state.gen_fn, param_name)
-        grad = get_param_grad(state.gen_fn, param_name)
-        set_param!(state.gen_fn, param_name, value + grad * step_size)
-        zero_param_grad!(state.gen_fn, param_name)
-    end
-    state.t += 1
-end

src/dynamic/backprop.jl

@@ -42,10 +42,10 @@ mutable struct GFBackpropParamsState
     visitor::AddressVisitor
     scale_factor::Float64
     active_gen_fn::GenerativeFunction
-    tracked_params::Dict{ParameterID,Any}
+    tracked_params::Dict{Tuple{GenerativeFunction,Symbol},Any}
 
     function GFBackpropParamsState(trace::DynamicDSLTrace, tape, scale_factor)
-        tracked_params = Dict{ParameterID,Any}()
+        tracked_params = Dict{Tuple{GenerativeFunction,Symbol},Any}()
         store = get_parameter_store(trace)
         gen_fn = get_gen_fn(trace)
         for (name, value) in get_local_parameters(store, gen_fn)
@@ -178,9 +178,10 @@ function accumulate_param_gradients!(trace::DynamicDSLTrace, retval_grad, scale_
     reverse_pass!(tape)
 
     # increment the gradient accumulators for trainable parameters in scope
+    store = get_parameter_store(trace)
     for ((active_gen_fn, name), tracked) in state.tracked_params
-        parameter_id = (active_gen_fn, parameter_id)
-	    increment_gradient!(store, parameter_id, deriv(tracked), state.scale_factor)
+        parameter_id = (active_gen_fn, name)
+	    increment_gradient!(parameter_id, deriv(tracked), state.scale_factor, store)
     end
 
     # return gradients with respect to arguments with gradients, or nothing
@@ -217,6 +218,16 @@ function GFBackpropTraceState(trace, selection, tape)
         get_gen_fn(trace))
 end
 
+get_parameter_store(state::GFBackpropTraceState) = get_parameter_store(state.trace)
+
+get_parameter_id(state::GFBackpropTraceState, name::Symbol) = (state.active_gen_fn, name)
+
+get_active_gen_fn(state::GFBackpropTraceState) = state.active_gen_fn
+
+function set_active_gen_fn!(state::GFBackpropTraceState, gen_fn::GenerativeFunction)
+    state.active_gen_fn = gen_fn
+end
+
 function fill_submaps!(
         map::DynamicChoiceMap,
         tracked_trie::Trie{Any,Union{TrackedReal,TrackedArray}},
@@ -304,15 +315,6 @@ function traceat(state::GFBackpropTraceState, gen_fn::GenerativeFunction{T,U},
     retval_maybe_tracked
 end
 
-function splice(state::GFBackpropTraceState, gen_fn::DynamicDSLFunction,
-                args_maybe_tracked::Tuple)
-    prev_gen_fn = state.active_gen_fn
-    state.active_gen_fn = gen_fn
-    retval = exec(gen_fn, state, args)
-    state.active_gen_fn = prev_gen_fn
-    return retval
-end
-
 @noinline function ReverseDiff.special_reverse_exec!(
         instruction::ReverseDiff.SpecialInstruction{BackpropTraceRecord})
     record = instruction.func

src/dynamic/dynamic.jl

@@ -29,14 +29,22 @@ function DynamicDSLFunction(arg_types::Vector{Type},
                 has_argument_grads, accepts_output_grad)
 end
 
-function DynamicDSLTrace(gen_fn::T, args) where {T<:DynamicDSLFunction}
+function Base.show(io::IO, gen_fn::DynamicDSLFunction)
+    return "Gen DML generative function: $(gen_fn.julia_function)"
+end
+
+function Base.show(io::IO, ::MIME"text/plain", gen_fn::DynamicDSLFunction)
+    return "Gen DML generative function: $(gen_fn.julia_function)"
+end
+
+function DynamicDSLTrace(gen_fn::T, args, parameter_store::JuliaParameterStore) where {T<:DynamicDSLFunction}
     # pad args with default values, if available
     if gen_fn.has_defaults && length(args) < length(gen_fn.arg_defaults)
         defaults = gen_fn.arg_defaults[length(args)+1:end]
         defaults = map(x -> something(x), defaults)
         args = Tuple(vcat(collect(args), defaults))
     end
-    return DynamicDSLTrace{T}(gen_fn, args)
+    return DynamicDSLTrace{T}(gen_fn, args, parameter_store)
 end
 
 accepts_output_grad(gen_fn::DynamicDSLFunction) = gen_fn.accepts_output_grad
@@ -58,6 +66,14 @@ function exec(gen_fn::DynamicDSLFunction, state, args::Tuple)
     gen_fn.julia_function(state, args...)
 end
 
+function splice(state, gen_fn::DynamicDSLFunction, args::Tuple)
+    prev_gen_fn = get_active_gen_fn(state)
+    state.active_gen_fn = gen_fn
+    retval = exec(gen_fn, state, args)
+    set_active_gen_fn!(state, prev_gen_fn)
+    return retval
+end
+
 # whether there is a gradient of score with respect to each argument
 # it returns 'nothing' for those arguemnts that don't have a derivatice
 has_argument_grads(gen::DynamicDSLFunction) = gen.has_argument_grads
@@ -104,7 +120,7 @@ end
 function read_param(state, name::Symbol)
     parameter_id = get_parameter_id(state, name)
     store = get_parameter_store(state)
-    return get_parameter_value(store, parameter_id)
+    return get_parameter_value(parameter_id, store)
 end
 
 ##################