Merge pull request #324 from probcomp/update-no-change

Add shorthand variants of `update` and `regenerate` that assume the arguments are unchanged

Author: bzinberg

src/gen_fn_interface.jl

@@ -276,6 +276,19 @@ function update(trace, args::Tuple, argdiffs::Tuple, ::ChoiceMap)
     error("Not implemented")
 end
 
+"""
+    (new_trace, weight, retdiff, discard) = update(trace, constraints::ChoiceMap)
+
+Shorthand variant of
+[`update`](@ref update(::Any, ::Tuple, ::Tuple, ::ChoiceMap))
+which assumes the arguments are unchanged.
+"""
+function update(trace, constraints::ChoiceMap)
+    args = get_args(trace)
+    argdiffs = Tuple(NoChange() for _ in args)
+    return update(trace, args, argdiffs, constraints)
+end
+
 """
     (new_trace, weight, retdiff) = regenerate(trace, args::Tuple, argdiffs::Tuple,
                                               selection::Selection)
@@ -307,6 +320,19 @@ function regenerate(trace, args::Tuple, argdiffs::Tuple, selection::Selection)
     error("Not implemented")
 end
 
+"""
+    (new_trace, weight, retdiff) = regenerate(trace, selection::Selection)
+
+Shorthand variant of
+[`regenerate`](@ref regenerate(::Any, ::Tuple, ::Tuple, ::Selection))
+which assumes the arguments are unchanged.
+"""
+function regenerate(trace, selection::Selection)
+    args = get_args(trace)
+    argdiffs = Tuple(NoChange() for _ in args)
+    return regenerate(trace, args, argdiffs, selection)
+end
+
 """
     arg_grads = accumulate_param_gradients!(trace, retgrad=nothing, scale_factor=1.)
 

test/gen_fn_interface.jl

@@ -0,0 +1,35 @@
+@gen function f_dynamic(x, y)
+    z ~ normal(0, 1)
+    return x + y + z
+end
+@gen (static) function f_static(x, y)
+    z ~ normal(0, 1)
+    return x + y + z
+end
+@load_generated_functions()
+
+for (lang, f) in [:dynamic => f_dynamic,
+                  :static  => f_static]
+    @testset "update(...) shorthand assuming unchanged args ($lang modeling lang)" begin
+        trace0 = simulate(f, (5, 6))
+
+        constraints = choicemap((:z, 0))
+        trace1, _, _, discard = update(trace0, constraints)
+        # The main test is that the shorthand version runs without crashing,
+        # which is already shown by the time we get here.  Beyond that, let's
+        # sanity-check that `update` did what it's supposed to.
+        @test get_args(trace1) == (5, 6)
+        @test trace1[:z] == 0
+        @test :z in keys(get_values_shallow(discard))
+    end
+
+    @testset "regenerate(...) shorthand assuming unchanged args ($lang modeling lang)" begin
+        trace0 = simulate(f, (5, 6))
+        trace1, _, _ = regenerate(trace0, select(:z))
+        # The main test is that the shorthand version runs without crashing,
+        # which is already shown by the time we get here.  Beyond that, let's
+        # sanity-check that `regenerate` did what it's supposed to.
+        @test get_args(trace1) == (5, 6)
+        @test trace1[:z] != trace0[:z]
+    end
+end

test/runtests.jl

@@ -68,6 +68,7 @@ include("autodiff.jl")
 include("diff.jl")
 include("selection.jl")
 include("assignment.jl")
+include("gen_fn_interface.jl")
 include("dsl/dsl.jl")
 include("optional_args.jl")
 include("static_ir/static_ir.jl")