Merge pull request #54 from TuringLang/dottilde

Take @. seriously and avoid eval

Author: devmotion

src/compiler.jl

@@ -246,40 +246,40 @@ function replace_sampler!(model_info)
     return model_info
 end
 
-# The next function is defined that way because .~ gives a parsing error in Julia 1.0
 """
-\"""
     replace_tilde!(model_info)
 
-Replaces `~` expressions with observation or assumption expressions, updating `model_info`.
-\"""
+Replace `~` and `.~` expressions with observation or assumption expressions, updating `model_info`.
+"""
 function replace_tilde!(model_info)
-    ex = model_info[:main_body]
-    ex = MacroTools.postwalk(ex) do x 
-        if @capture(x, @M_ L_ ~ R_) && M == Symbol("@__dot__")
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
-        end
-    end
-    $(VERSION >= v"1.1" ? "ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ .~ R_)
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
+    # Apply the `@.` macro first.
+    expr = model_info[:main_body]
+    dottedexpr = MacroTools.postwalk(apply_dotted, expr)
+
+    # Check for tilde operators.
+    tildeexpr = MacroTools.postwalk(dottedexpr) do x
+        # Check dot tilde first.
+        dotargs = getargs_dottilde(x)
+        if dotargs !== nothing
+            L, R = dotargs
+            return generate_dot_tilde(L, R, model_info)
         end
-    end" : "")
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ ~ R_)
-            generate_tilde(L, R, model_info)
-        else
-            x
+
+        # Check tilde.
+        args = getargs_tilde(x)
+        if args !== nothing
+            L, R = args
+            return generate_tilde(L, R, model_info)
         end
+
+        return x
     end
-    model_info[:main_body] = ex
+
+    # Update the function body.
+    model_info[:main_body] = tildeexpr
+
     return model_info
 end
-""" |> Meta.parse |> eval
 
 # """ Unbreak code highlighting in Emacs julia-mode
 

src/utils.jl

@@ -1,3 +1,53 @@
+"""
+    apply_dotted(x)
+
+Apply the transformation of the `@.` macro if `x` is an expression of the form `@. X`.
+"""
+apply_dotted(x) = x
+function apply_dotted(expr::Expr)
+    if Meta.isexpr(expr, :macrocall) && length(expr.args) > 1 &&
+        expr.args[1] === Symbol("@__dot__")
+        return Base.Broadcast.__dot__(expr.args[end])
+    end
+    return expr
+end
+
+"""
+    getargs_dottilde(x)
+
+Return the arguments `L` and `R`, if `x` is an expression of the form `L .~ R` or
+`(~).(L, R)`, or `nothing` otherwise.
+"""
+getargs_dottilde(x) = nothing
+function getargs_dottilde(expr::Expr)
+    # Check if the expression is of the form `L .~ R`.
+    if Meta.isexpr(expr, :call, 3) && expr.args[1] === :.~
+        return expr.args[2], expr.args[3]
+    end
+
+    # Check if the expression is of the form `(~).(L, R)`.
+    if Meta.isexpr(expr, :., 2) && expr.args[1] === :~ &&
+        Meta.isexpr(expr.args[2], :tuple, 2)
+        return expr.args[2].args[1], expr.args[2].args[2]
+    end
+
+    return
+end
+
+"""
+    getargs_tilde(x)
+
+Return the arguments `L` and `R`, if `x` is an expression of the form `L ~ R`, or `nothing`
+otherwise.
+"""
+getargs_tilde(x) = nothing
+function getargs_tilde(expr::Expr)
+    if Meta.isexpr(expr, :call, 3) && expr.args[1] === :~
+        return expr.args[2], expr.args[3]
+    end
+    return
+end
+
 ############################################
 # Julia 1.2 temporary fix - Julia PR 33303 #
 ############################################

test/compiler.jl

@@ -327,14 +327,23 @@ priors = 0 # See "new grammar" test.
         x = randn(100)
         res = sample(vdemo1(x), alg, 250)
 
+        @model vdemo1b(x) = begin
+            s ~ InverseGamma(2,3)
+            m ~ Normal(0, sqrt(s))
+            @. x ~ Normal(m, $(sqrt(s)))
+            return s, m
+        end
+
+        res = sample(vdemo1b(x), alg, 250)
+
         D = 2
         @model vdemo2(x) = begin
             μ ~ MvNormal(zeros(D), ones(D))
-            @. x ~ MvNormal(μ, ones(D))
+            @. x ~ $(MvNormal(μ, ones(D)))
         end
 
         alg = HMC(0.01, 5)
-        res = sample(vdemo2(randn(D,100)), alg, 250)
+        res = sample(vdemo2(randn(D, 100)), alg, 250)
 
         # Vector assumptions
         N = 10
@@ -386,78 +395,75 @@ priors = 0 # See "new grammar" test.
         sample(vdemo7(), alg, 1000)
     end
 
-    if VERSION >= v"1.1"
-        """
-        @testset "vectorization .~" begin
-            @model vdemo1(x) = begin
-                s ~ InverseGamma(2,3)
-                m ~ Normal(0, sqrt(s))
-                x .~ Normal(m, sqrt(s))
-                return s, m
-            end
+    # Notation is ugly since `x .~ Normal(μ, σ)` cannot be parsed in Julia 1.0
+    @testset "vectorization .~" begin
+        @model vdemo1(x) = begin
+            s ~ InverseGamma(2,3)
+            m ~ Normal(0, sqrt(s))
+            (.~)(x, Normal(m, sqrt(s)))
+            return s, m
+        end
 
-            alg = HMC(0.01, 5)
-            x = randn(100)
-            res = sample(vdemo1(x), alg, 250)
+        alg = HMC(0.01, 5)
+        x = randn(100)
+        res = sample(vdemo1(x), alg, 250)
 
-            D = 2
-            @model vdemo2(x) = begin
-                μ ~ MvNormal(zeros(D), ones(D))
-                x .~ MvNormal(μ, ones(D))
-            end
+        D = 2
+        @model vdemo2(x) = begin
+            μ ~ MvNormal(zeros(D), ones(D))
+            (.~)(x, MvNormal(μ, ones(D)))
+        end
 
-            alg = HMC(0.01, 5)
-            res = sample(vdemo2(randn(D,100)), alg, 250)
+        alg = HMC(0.01, 5)
+        res = sample(vdemo2(randn(D,100)), alg, 250)
 
-            # Vector assumptions
-            N = 10
-            setchunksize(N)
-            alg = HMC(0.2, 4)
+        # Vector assumptions
+        N = 10
+        setchunksize(N)
+        alg = HMC(0.2, 4)
 
-            @model vdemo3() = begin
-                x = Vector{Real}(undef, N)
-                for i = 1:N
-                    x[i] ~ Normal(0, sqrt(4))
-                end
+        @model vdemo3() = begin
+            x = Vector{Real}(undef, N)
+            for i = 1:N
+                x[i] ~ Normal(0, sqrt(4))
             end
+        end
 
-            t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
+        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
 
-            # Test for vectorize UnivariateDistribution
-            @model vdemo4() = begin
+        # Test for vectorize UnivariateDistribution
+        @model vdemo4() = begin
             x = Vector{Real}(undef, N)
-            x .~ Normal(0, 2)
-            end
+            (.~)(x, Normal(0, 2))
+        end
 
-            t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
+        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
 
-            @model vdemo5() = begin
-                x ~ MvNormal(zeros(N), 2 * ones(N))
-            end
+        @model vdemo5() = begin
+            x ~ MvNormal(zeros(N), 2 * ones(N))
+        end
 
-            t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
+        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
 
-            println("Time for")
-            println("  Loop : \$t_loop")
-            println("  Vec  : \$t_vec")
-            println("  Mv   : \$t_mv")
+        println("Time for")
+        println("  Loop : \$t_loop")
+        println("  Vec  : \$t_vec")
+        println("  Mv   : \$t_mv")
 
-            # Transformed test
-            @model vdemo6() = begin
-                x = Vector{Real}(undef, N)
-                x .~ InverseGamma(2, 3)
-            end
+        # Transformed test
+        @model vdemo6() = begin
+            x = Vector{Real}(undef, N)
+            (.~)(x, InverseGamma(2, 3))
+        end
 
-            sample(vdemo6(), alg, 1000)
+        sample(vdemo6(), alg, 1000)
 
-            @model vdemo7() = begin
-                x = Array{Real}(undef, N, N)
-                x .~ [InverseGamma(2, 3) for i in 1:N]
-            end
-    
-            sample(vdemo7(), alg, 1000)
+        @model vdemo7() = begin
+            x = Array{Real}(undef, N, N)
+            (.~)(x, [InverseGamma(2, 3) for i in 1:N])
         end
-        """ |> Meta.parse |> eval
+
+        sample(vdemo7(), alg, 1000)
     end
 
     @testset "Type parameters" begin

test/runtests.jl

@@ -6,6 +6,7 @@ using .Turing
 turnprogress(false)
 
 @testset "DynamicPPL.jl" begin
+    include("utils.jl")
     include("compiler.jl")
     include("varinfo.jl")
     include("prob_macro.jl")

test/utils.jl

@@ -0,0 +1,50 @@
+using DynamicPPL
+using DynamicPPL: apply_dotted, getargs_dottilde, getargs_tilde
+
+using Test
+
+@testset "apply_dotted" begin
+    # Some things that are not expressions.
+    @test apply_dotted(:x) === :x
+    @test apply_dotted(1.0) === 1.0
+    @test apply_dotted([1.0, 2.0, 4.0]) == [1.0, 2.0, 4.0]
+
+    # Some expressions.
+    @test apply_dotted(:(x ~ Normal(μ, σ))) == :(x ~ Normal(μ, σ))
+    @test apply_dotted(:((.~)(x, Normal(μ, σ)))) == :((.~)(x, Normal(μ, σ)))
+    @test apply_dotted(:((~).(x, Normal(μ, σ)))) == :((~).(x, Normal(μ, σ)))
+    @test apply_dotted(:(@. x ~ Normal(μ, σ))) == :((~).(x, Normal.(μ, σ)))
+    @test apply_dotted(:(@. x ~ Normal(μ, $(Expr(:$, :(sqrt(v))))))) ==
+        :((~).(x, Normal.(μ, sqrt(v))))
+    @test apply_dotted(:(@~ Normal.(μ, σ))) == :(@~ Normal.(μ, σ))
+end
+
+@testset "getargs_dottilde" begin
+    # Some things that are not expressions.
+    @test getargs_dottilde(:x) === nothing
+    @test getargs_dottilde(1.0) === nothing
+    @test getargs_dottilde([1.0, 2.0, 4.0]) === nothing
+
+    # Some expressions.
+    @test getargs_dottilde(:(x ~ Normal(μ, σ))) === nothing
+    @test getargs_dottilde(:((.~)(x, Normal(μ, σ)))) == (:x, :(Normal(μ, σ)))
+    @test getargs_dottilde(:((~).(x, Normal(μ, σ)))) == (:x, :(Normal(μ, σ)))
+    @test getargs_dottilde(:(@. x ~ Normal(μ, σ))) === nothing
+    @test getargs_dottilde(:(@. x ~ Normal(μ, $(Expr(:$, :(sqrt(v))))))) === nothing
+    @test getargs_dottilde(:(@~ Normal.(μ, σ))) === nothing
+end
+
+@testset "getargs_tilde" begin
+    # Some things that are not expressions.
+    @test getargs_tilde(:x) === nothing
+    @test getargs_tilde(1.0) === nothing
+    @test getargs_tilde([1.0, 2.0, 4.0]) === nothing
+
+    # Some expressions.
+    @test getargs_tilde(:(x ~ Normal(μ, σ))) == (:x, :(Normal(μ, σ)))
+    @test getargs_tilde(:((.~)(x, Normal(μ, σ)))) === nothing
+    @test getargs_tilde(:((~).(x, Normal(μ, σ)))) === nothing
+    @test getargs_tilde(:(@. x ~ Normal(μ, σ))) === nothing
+    @test getargs_tilde(:(@. x ~ Normal(μ, $(Expr(:$, :(sqrt(v))))))) === nothing
+    @test getargs_tilde(:(@~ Normal.(μ, σ))) === nothing
+end