Remove in-place updates and fix return values (#64)

* Remove in-place updates and fix return values

* Bump version

* Fix format

* added tests for all ADs

* test - just sample 2 observations single-threaded

* fix format with empty line

* remove rdcache

* updating jose's email in Project.toml

* Update Project.toml

* oops @Assert is actually @test

* see if tests passes without Tracker

* removed zygote also to see if tests passes

Co-authored-by: Jose Storopoli <[email protected]>

Author: devmotion

Project.toml

@@ -1,7 +1,8 @@
 name = "TuringGLM"
 uuid = "0004c1f4-53c5-4d43-a221-a1dac6cf6b74"
-authors = ["Jose Storopoli <[email protected]>, Rik Huijzer <[email protected]>, and contributors"]
-version = "2.1.1"
+authors = ["Jose Storopoli <[email protected]>, Rik Huijzer <[email protected]>, and contributors"]
+version = "2.1.2"
+
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

src/turing_model.jl

@@ -183,16 +183,16 @@ function _model(μ_X, σ_X, prior, intercept_ranef, idx, ::Type{Normal})
         α ~ prior.intercept
         β ~ filldist(prior.predictors, predictors)
         σ ~ Exponential(residual)
-        μ = α .+ X * β
-        if !isempty(intercept_ranef)
+        if isempty(intercept_ranef)
+            μ = α .+ X * β
+        else
             τ ~ mad_y * truncated(TDist(3); lower=0)
             zⱼ ~ filldist(Normal(), n_gr)
-            αⱼ = zⱼ .* τ
-            μ .+= αⱼ[idxs]
+            μ = α .+ τ .* getindex.((zⱼ,), idxs) .+ X * β
         end
         #TODO: implement random-effects slope
         y ~ MvNormal(μ, σ^2 * I)
-        return (; α, β, σ, τ, zⱼ, αⱼ, y)
+        return nothing
     end
 end
 function _model(μ_X, σ_X, prior, ::Type{Normal})
@@ -203,7 +203,7 @@ function _model(μ_X, σ_X, prior, ::Type{Normal})
         β ~ filldist(prior.predictors, predictors)
         σ ~ Exponential(residual)
         y ~ MvNormal(α .+ X * β, σ^2 * I)
-        return (; α, β, σ, y)
+        return nothing
     end
 end
 
@@ -226,16 +226,16 @@ function _model(μ_X, σ_X, prior, intercept_ranef, idx, ::Type{TDist})
         β ~ filldist(prior.predictors, predictors)
         σ ~ Exponential(residual)
         ν ~ prior.auxiliary
-        μ = α .+ X * β
-        if !isempty(intercept_ranef)
+        if isempty(intercept_ranef)
+            μ = α .+ X * β
+        else
             τ ~ mad_y * truncated(TDist(3); lower=0)
             zⱼ ~ filldist(Normal(), n_gr)
-            αⱼ = zⱼ .* τ
-            μ .+= αⱼ[idxs]
+            μ = α .+ τ .* getindex.((zⱼ,), idxs) .+ X * β
         end
         #TODO: implement random-effects slope
         y ~ arraydist(μ + σ * TDist.(ν))
-        return (; α, β, σ, ν, τ, zⱼ, αⱼ, y)
+        return nothing
     end
 end
 function _model(μ_X, σ_X, prior, ::Type{TDist})
@@ -247,7 +247,7 @@ function _model(μ_X, σ_X, prior, ::Type{TDist})
         σ ~ Exponential(residual)
         ν ~ prior.auxiliary
         y ~ arraydist((α .+ X * β) .+ σ .* TDist.(ν))
-        return (; α, β, σ, ν, y)
+        return nothing
     end
 end
 
@@ -267,16 +267,16 @@ function _model(μ_X, σ_X, prior, intercept_ranef, idx, ::Type{Bernoulli})
     )
         α ~ prior.intercept
         β ~ filldist(prior.predictors, predictors)
-        μ = α .+ X * β
-        if !isempty(intercept_ranef)
+        if isempty(intercept_ranef)
+            μ = α .+ X * β
+        else
             τ ~ mad_y * truncated(TDist(3); lower=0)
             zⱼ ~ filldist(Normal(), n_gr)
-            αⱼ = zⱼ .* τ
-            μ .+= αⱼ[idxs]
+            μ = α .+ τ .* getindex.((zⱼ,), idxs) .+ X * β
         end
         #TODO: implement random-effects slope
         y ~ arraydist(LazyArray(@~ BernoulliLogit.(μ)))
-        return (; α, β, τ, zⱼ, αⱼ, y)
+        return nothing
     end
 end
 function _model(μ_X, σ_X, prior, ::Type{Bernoulli})
@@ -286,7 +286,7 @@ function _model(μ_X, σ_X, prior, ::Type{Bernoulli})
         α ~ prior.intercept
         β ~ filldist(prior.predictors, predictors)
         y ~ arraydist(LazyArray(@~ BernoulliLogit.(α .+ X * β)))
-        return (; α, β, y)
+        return nothing
     end
 end
 
@@ -306,16 +306,16 @@ function _model(μ_X, σ_X, prior, intercept_ranef, idx, ::Type{Poisson})
     )
         α ~ prior.intercept
         β ~ filldist(prior.predictors, predictors)
-        μ = α .+ X * β
-        if !isempty(intercept_ranef)
+        if isempty(intercept_ranef)
+            μ = α .+ X * β
+        else
             τ ~ mad_y * truncated(TDist(3); lower=0)
             zⱼ ~ filldist(Normal(), n_gr)
-            αⱼ = zⱼ .* τ
-            μ .+= αⱼ[idxs]
+            μ = α .+ τ .* getindex.((zⱼ,), idxs) .+ X * β
         end
         #TODO: implement random-effects slope
         y ~ arraydist(LazyArray(@~ LogPoisson.(μ)))
-        return (; α, β, τ, zⱼ, αⱼ, y)
+        return nothing
     end
 end
 function _model(μ_X, σ_X, prior, ::Type{Poisson})
@@ -325,7 +325,7 @@ function _model(μ_X, σ_X, prior, ::Type{Poisson})
         α ~ prior.intercept
         β ~ filldist(prior.predictors, predictors)
         y ~ arraydist(LazyArray(@~ LogPoisson.(α .+ X * β)))
-        return (; α, β, y)
+        return nothing
     end
 end
 
@@ -347,16 +347,16 @@ function _model(μ_X, σ_X, prior, intercept_ranef, idx, ::Type{NegativeBinomial
         β ~ filldist(prior.predictors, predictors)
         ϕ⁻ ~ prior.auxiliary
         ϕ = 1 / ϕ⁻
-        μ = α .+ X * β
-        if !isempty(intercept_ranef)
+        if isempty(intercept_ranef)
+            μ = α .+ X * β
+        else
             τ ~ mad_y * truncated(TDist(3); lower=0)
             zⱼ ~ filldist(Normal(), n_gr)
-            αⱼ = zⱼ .* τ
-            μ .+= αⱼ[idxs]
+            μ = α .+ τ .* getindex.((zⱼ,), idxs) .+ X * β
         end
         #TODO: implement random-effects slope
         y ~ arraydist(LazyArray(@~ NegativeBinomial2.(exp.(μ), ϕ)))
-        return (; α, β, ϕ, τ, zⱼ, αⱼ, y)
+        return nothing
     end
 end
 function _model(μ_X, σ_X, prior, ::Type{NegativeBinomial})
@@ -366,7 +366,7 @@ function _model(μ_X, σ_X, prior, ::Type{NegativeBinomial})
         ϕ⁻ ~ prior.auxiliary
         ϕ = 1 / ϕ⁻
         y ~ arraydist(LazyArray(@~ NegativeBinomial2.(exp.(α .+ X * β), ϕ)))
-        return (; α, β, ϕ, y)
+        return nothing
     end
 end
 

test/Project.toml

@@ -2,10 +2,14 @@
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 CSV = "0.9, 1"

test/ad_backends.jl

@@ -0,0 +1,34 @@
+@timed_testset "ad_backends" begin
+    DATA_DIR = joinpath("..", "data")
+    cheese = CSV.read(joinpath(DATA_DIR, "cheese.csv"), DataFrame)
+    f = @formula(y ~ (1 | cheese) + background)
+    m = turing_model(f, cheese)
+    # only running 2 samples to test if the different ADs runs
+    @timed_testset "ForwardDiff" begin
+        Turing.setadbackend(:forwarddiff)
+        chn = sample(m, NUTS(), 2)
+        @test chn isa Chains
+    end
+    # TODO: fix Tracker tests
+    # @timed_testset "Tracker" begin
+    #     using Tracker
+    #     Turing.setadbackend(:tracker)
+    #     chn = sample(m, NUTS(), 2)
+    #     @test chn isa Chains
+    # end
+    # TODO: fix Zygote tests
+    # @timed_testset "Zygote" begin
+    #     using Zygote
+    #     Turing.setadbackend(:zygote)
+    #     chn = sample(m, NUTS(), 2)
+    #     @test chn isa Chains
+    # end
+    @timed_testset "ReverseDiff" begin
+        using ReverseDiff
+        Turing.setadbackend(:reversediff)
+        chn = sample(m, NUTS(), 2)
+        @test chn isa Chains
+    end
+    # go back to defaults
+    Turing.setadbackend(:forwarddiff)
+end

test/runtests.jl

@@ -50,6 +50,7 @@ end
     include("utils.jl")
     include("priors.jl")
     include("turing_model.jl")
+    include("ad_backends.jl")
 end
 
 show(TIMEROUTPUT; compact=true, sortby=:firstexec)