Revert "Revert today's commits"

penelopeysm · penelopeysm · commit 1f4206db48c5 · 2025-05-19T13:22:36.000+01:00
This reverts commit 78a2913.
diff --git a/.github/workflows/generate_website.yml b/.github/workflows/generate_website.yml
@@ -7,14 +7,16 @@ on:
     # If _any_ of the modified files match this filter, it will trigger this
     # workflow
     paths:
-      - '*.jl'
+      - '**/*.jl'
       - '*.toml'
       - 'ad.py'
+      - '.github/workflows/generate_website.yml'
   pull_request:
     paths:
-      - '*.jl'
+      - '**/*.jl'
       - '*.toml'
       - 'ad.py'
+      - '.github/workflows/generate_website.yml'
   workflow_dispatch:
 
 permissions:
diff --git a/.github/workflows/refresh_website.yml b/.github/workflows/refresh_website.yml
@@ -7,14 +7,16 @@ on:
     # If _all_ the modified files match this filter, it won't trigger this
     # workflow
     paths-ignore:
-      - '*.jl'
+      - '**/*.jl'
       - '*.toml'
       - 'ad.py'
+      - '.github/workflows/refresh_website.yml'
   pull_request:
     paths:
-      - '*.jl'
+      - '**/*.jl'
       - '*.toml'
       - 'ad.py'
+      - '.github/workflows/refresh_website.yml'
   workflow_dispatch:
 
 permissions:
diff --git a/ad.py b/ad.py
@@ -89,8 +89,12 @@ def run_ad(args):
         try:
             output = run_and_capture([*RUN_JULIA_COMMAND, "--run", model_key, adtype])
             result = try_float(output.splitlines()[-1])
+            if not isinstance(result, float):
+                print(f"Output: {output}")
         except sp.CalledProcessError as e:
-            print(f"Error running {model_key} with {adtype}. Output: {e.output}")
+            # Julia crashed
+            print(f"Julia crashed when running {model_key} with {adtype}.")
+            print(f"To reproduce, run: `julia --project=. main.jl --run {model_key} {adtype}`")
             result = "error"
 
         print(f" ... {model_key} with {adtype} ==> {result}")
diff --git a/main.jl b/main.jl
@@ -13,13 +13,13 @@ import Zygote
 
 # AD backends to test.
 ADTYPES = Dict(
-    "FiniteDifferences" => AutoFiniteDifferences(; fdm=central_fdm(5, 1)),
+    "FiniteDifferences" => AutoFiniteDifferences(; fdm = central_fdm(5, 1)),
     "ForwardDiff" => AutoForwardDiff(),
-    "ReverseDiff" => AutoReverseDiff(; compile=false),
-    "ReverseDiffCompiled" => AutoReverseDiff(; compile=true),
-    "Mooncake" => AutoMooncake(; config=nothing),
-    "EnzymeForward" => AutoEnzyme(; mode=set_runtime_activity(Forward, true)),
-    "EnzymeReverse" => AutoEnzyme(; mode=set_runtime_activity(Reverse, true)),
+    "ReverseDiff" => AutoReverseDiff(; compile = false),
+    "ReverseDiffCompiled" => AutoReverseDiff(; compile = true),
+    "Mooncake" => AutoMooncake(; config = nothing),
+    "EnzymeForward" => AutoEnzyme(; mode = set_runtime_activity(Forward, true)),
+    "EnzymeReverse" => AutoEnzyme(; mode = set_runtime_activity(Reverse, true)),
     "Zygote" => AutoZygote(),
 )
 
@@ -35,7 +35,7 @@ end
 # These imports tend to get used a lot in models
 using DynamicPPL: @model, to_submodel
 using Distributions
-using LinearAlgebra: I
+using LinearAlgebra
 
 include("models/assume_dirichlet.jl")
 include("models/assume_lkjcholu.jl")
@@ -44,8 +44,21 @@ include("models/assume_normal.jl")
 include("models/assume_submodel.jl")
 include("models/assume_wishart.jl")
 include("models/control_flow.jl")
-include("models/dot_assume_observe_index.jl")
+include("models/demo_assume_dot_observe_literal.jl")
+include("models/demo_assume_dot_observe.jl")
+include("models/demo_assume_index_observe.jl")
+include("models/demo_assume_matrix_observe_matrix_index.jl")
+include("models/demo_assume_multivariate_observe_literal.jl")
+include("models/demo_assume_multivariate_observe.jl")
+include("models/demo_assume_observe_literal.jl")
+include("models/demo_assume_submodel_observe_index_literal.jl")
+include("models/demo_dot_assume_observe_index_literal.jl")
+include("models/demo_dot_assume_observe_index.jl")
+include("models/demo_dot_assume_observe_matrix_index.jl")
+include("models/demo_dot_assume_observe_submodel.jl")
+include("models/demo_dot_assume_observe.jl")
 include("models/dot_assume.jl")
+include("models/demo_dot_assume_observe.jl")
 include("models/dot_observe.jl")
 include("models/dynamic_constraint.jl")
 include("models/multiple_constraints_same_var.jl")
@@ -74,9 +87,9 @@ elseif length(ARGS) == 3 && ARGS[1] == "--run"
             # https://github.com/TuringLang/ADTests/issues/4
             vi = DynamicPPL.unflatten(VarInfo(model), [0.5, -0.5])
             params = [-0.5, 0.5]
-            result = run_ad(model, adtype; varinfo=vi, params=params, benchmark=true)
+            result = run_ad(model, adtype; varinfo = vi, params = params, benchmark = true)
         else
-            result = run_ad(model, adtype; benchmark=true)
+            result = run_ad(model, adtype; benchmark = true)
         end
         # If reached here - nothing went wrong
         @printf("%.3f", result.time_vs_primal)
diff --git a/models/demo_assume_dot_observe.jl b/models/demo_assume_dot_observe.jl
@@ -0,0 +1,8 @@
+@model function demo_assume_dot_observe(x = [1.5, 2.0])
+    # `assume` and `dot_observe`
+    s ~ InverseGamma(2, 3)
+    m ~ Normal(0, sqrt(s))
+    x .~ Normal(m, sqrt(s))
+end
+
+@register demo_assume_dot_observe()
diff --git a/models/demo_assume_dot_observe_literal.jl b/models/demo_assume_dot_observe_literal.jl
@@ -0,0 +1,8 @@
+@model function demo_assume_dot_observe_literal()
+    # `assume` and literal `dot_observe`
+    s ~ InverseGamma(2, 3)
+    m ~ Normal(0, sqrt(s))
+    [1.5, 2.0] .~ Normal(m, sqrt(s))
+end
+
+@register demo_assume_dot_observe_literal()
diff --git a/models/demo_assume_index_observe.jl b/models/demo_assume_index_observe.jl
@@ -0,0 +1,17 @@
+@model function demo_assume_index_observe(
+    x = [1.5, 2.0],
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    # `assume` with indexing and `observe`
+    s = TV(undef, length(x))
+    for i in eachindex(s)
+        s[i] ~ InverseGamma(2, 3)
+    end
+    m = TV(undef, length(x))
+    for i in eachindex(m)
+        m[i] ~ Normal(0, sqrt(s[i]))
+    end
+    x ~ MvNormal(m, Diagonal(s))
+end
+
+@register demo_assume_index_observe()
diff --git a/models/demo_assume_matrix_observe_matrix_index.jl b/models/demo_assume_matrix_observe_matrix_index.jl
@@ -0,0 +1,14 @@
+@model function demo_assume_matrix_observe_matrix_index(
+    x = transpose([1.5 2.0;]),
+    ::Type{TV} = Array{Float64},
+) where {TV}
+    n = length(x)
+    d = n ÷ 2
+    s ~ reshape(product_distribution(fill(InverseGamma(2, 3), n)), d, 2)
+    s_vec = vec(s)
+    m ~ MvNormal(zeros(n), Diagonal(s_vec))
+
+    x[:, 1] ~ MvNormal(m, Diagonal(s_vec))
+end
+
+@register demo_assume_matrix_observe_matrix_index()
diff --git a/models/demo_assume_multivariate_observe.jl b/models/demo_assume_multivariate_observe.jl
@@ -0,0 +1,7 @@
+@model function demo_assume_multivariate_observe(x = [1.5, 2.0])
+    # Multivariate `assume` and `observe`
+    s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)])
+    m ~ MvNormal(zero(x), Diagonal(s))
+    x ~ MvNormal(m, Diagonal(s))
+end
+@register demo_assume_multivariate_observe()
diff --git a/models/demo_assume_multivariate_observe_literal.jl b/models/demo_assume_multivariate_observe_literal.jl
@@ -0,0 +1,8 @@
+@model function demo_assume_multivariate_observe_literal()
+    # multivariate `assume` and literal `observe`
+    s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)])
+    m ~ MvNormal(zeros(2), Diagonal(s))
+    [1.5, 2.0] ~ MvNormal(m, Diagonal(s))
+end
+
+@register demo_assume_multivariate_observe_literal()
diff --git a/models/demo_assume_observe_literal.jl b/models/demo_assume_observe_literal.jl
@@ -0,0 +1,9 @@
+@model function demo_assume_observe_literal()
+    # univariate `assume` and literal `observe`
+    s ~ InverseGamma(2, 3)
+    m ~ Normal(0, sqrt(s))
+    1.5 ~ Normal(m, sqrt(s))
+    2.0 ~ Normal(m, sqrt(s))
+end
+
+@register demo_assume_observe_literal()
diff --git a/models/demo_assume_submodel_observe_index_literal.jl b/models/demo_assume_submodel_observe_index_literal.jl
@@ -0,0 +1,17 @@
+@model function _prior_dot_assume(::Type{TV} = Vector{Float64}) where {TV}
+    s = TV(undef, 2)
+    s .~ InverseGamma(2, 3)
+    m = TV(undef, 2)
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+    return s, m
+end
+
+@model function demo_assume_submodel_observe_index_literal()
+    # Submodel prior
+    priors ~ to_submodel(_prior_dot_assume(), false)
+    s, m = priors
+    1.5 ~ Normal(m[1], sqrt(s[1]))
+    2.0 ~ Normal(m[2], sqrt(s[2]))
+end
+
+@register demo_assume_submodel_observe_index_literal()
diff --git a/models/demo_dot_assume_observe.jl b/models/demo_dot_assume_observe.jl
@@ -0,0 +1,13 @@
+@model function demo_dot_assume_observe(
+    x = [1.5, 2.0],
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    # `dot_assume` and `observe`
+    s = TV(undef, length(x))
+    m = TV(undef, length(x))
+    s .~ InverseGamma(2, 3)
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+    x ~ MvNormal(m, Diagonal(s))
+end
+
+@register demo_dot_assume_observe()
diff --git a/models/demo_dot_assume_observe_index.jl b/models/demo_dot_assume_observe_index.jl
@@ -0,0 +1,15 @@
+@model function demo_dot_assume_observe_index(
+    x = [1.5, 2.0],
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    # `dot_assume` and `observe` with indexing
+    s = TV(undef, length(x))
+    s .~ InverseGamma(2, 3)
+    m = TV(undef, length(x))
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+    for i in eachindex(x)
+        x[i] ~ Normal(m[i], sqrt(s[i]))
+    end
+end
+
+@register demo_dot_assume_observe_index()
diff --git a/models/demo_dot_assume_observe_index_literal.jl b/models/demo_dot_assume_observe_index_literal.jl
@@ -0,0 +1,14 @@
+@model function demo_dot_assume_observe_index_literal(
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    # `dot_assume` and literal `observe` with indexing
+    s = TV(undef, 2)
+    m = TV(undef, 2)
+    s .~ InverseGamma(2, 3)
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+
+    1.5 ~ Normal(m[1], sqrt(s[1]))
+    2.0 ~ Normal(m[2], sqrt(s[2]))
+end
+
+@register demo_dot_assume_observe_index_literal()
diff --git a/models/demo_dot_assume_observe_matrix_index.jl b/models/demo_dot_assume_observe_matrix_index.jl
@@ -0,0 +1,12 @@
+@model function demo_dot_assume_observe_matrix_index(
+    x = transpose([1.5 2.0;]),
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    s = TV(undef, length(x))
+    s .~ InverseGamma(2, 3)
+    m = TV(undef, length(x))
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+    x[:, 1] ~ MvNormal(m, Diagonal(s))
+end
+
+@register demo_dot_assume_observe_matrix_index()
diff --git a/models/demo_dot_assume_observe_submodel.jl b/models/demo_dot_assume_observe_submodel.jl
@@ -0,0 +1,20 @@
+@model function _likelihood_multivariate_observe(s, m, x)
+    return x ~ MvNormal(m, Diagonal(s))
+end
+
+@model function demo_dot_assume_observe_submodel(
+    x = [1.5, 2.0],
+    ::Type{TV} = Vector{Float64},
+) where {TV}
+    s = TV(undef, length(x))
+    s .~ InverseGamma(2, 3)
+    m = TV(undef, length(x))
+    m ~ product_distribution(Normal.(0, sqrt.(s)))
+
+    # Submodel likelihood
+    # With to_submodel, we have to have a left-hand side variable to
+    # capture the result, so we just use a dummy variable
+    _ignore ~ to_submodel(_likelihood_multivariate_observe(s, m, x))
+end
+
+@register demo_dot_assume_observe_submodel()
diff --git a/models/dot_assume.jl b/models/dot_assume.jl
@@ -1,4 +1,4 @@
-@model function dot_assume(::Type{TV}=Vector{Float64}) where {TV}
+@model function dot_assume(::Type{TV} = Vector{Float64}) where {TV}
     a = TV(undef, 5)
     a .~ Normal()
 end
diff --git a/models/dot_assume_observe_index.jl b/models/dot_assume_observe_index.jl
diff --git a/models/dot_observe.jl b/models/dot_observe.jl
@@ -1,4 +1,4 @@
-@model function dot_observe(x=[1.5, 2.0, 2.5])
+@model function dot_observe(x = [1.5, 2.0, 2.5])
     a ~ Normal()
     x .~ Normal(a)
 end
diff --git a/models/dynamic_constraint.jl b/models/dynamic_constraint.jl
@@ -1,5 +1,6 @@
-@model function dynamic_constraint() a ~ Normal()
-    b ~ truncated(Normal(); lower=a)
+@model function dynamic_constraint()
+    a ~ Normal()
+    b ~ truncated(Normal(); lower = a)
 end
 
 @register dynamic_constraint()
diff --git a/models/multiple_constraints_same_var.jl b/models/multiple_constraints_same_var.jl
@@ -1,6 +1,4 @@
-@model function multiple_constraints_same_var(
-    ::Type{TV}=Vector{Float64}
-) where {TV}
+@model function multiple_constraints_same_var(::Type{TV} = Vector{Float64}) where {TV}
     x = TV(undef, 5)
     x[1] ~ Normal()
     x[2] ~ InverseGamma(2, 3)
diff --git a/models/n010.jl b/models/n010.jl
@@ -1,4 +1,4 @@
-@model function n010(::Type{TV}=Vector{Float64}) where {TV}
+@model function n010(::Type{TV} = Vector{Float64}) where {TV}
     a = TV(undef, 10)
     for i in eachindex(a)
         a[i] ~ Normal()
diff --git a/models/n050.jl b/models/n050.jl
@@ -1,4 +1,4 @@
-@model function n050(::Type{TV}=Vector{Float64}) where {TV}
+@model function n050(::Type{TV} = Vector{Float64}) where {TV}
     a = TV(undef, 50)
     for i in eachindex(a)
         a[i] ~ Normal()
diff --git a/models/n100.jl b/models/n100.jl
@@ -1,4 +1,4 @@
-@model function n100(::Type{TV}=Vector{Float64}) where {TV}
+@model function n100(::Type{TV} = Vector{Float64}) where {TV}
     a = TV(undef, 100)
     for i in eachindex(a)
         a[i] ~ Normal()
diff --git a/models/n500.jl b/models/n500.jl
@@ -1,4 +1,4 @@
-@model function n500(::Type{TV}=Vector{Float64}) where {TV}
+@model function n500(::Type{TV} = Vector{Float64}) where {TV}
     a = TV(undef, 500)
     for i in eachindex(a)
         a[i] ~ Normal()
diff --git a/models/observe_index.jl b/models/observe_index.jl
@@ -1,4 +1,4 @@
-@model function observe_index(x=[1.5, 2.0, 2.5])
+@model function observe_index(x = [1.5, 2.0, 2.5])
     a ~ Normal()
     for i in eachindex(x)
         x[i] ~ Normal(a)
diff --git a/models/observe_multivariate.jl b/models/observe_multivariate.jl
@@ -1,4 +1,7 @@
-@model function observe_multivariate(x=[1.5, 2.0, 2.5], ::Type{TV}=Vector{Float64}) where {TV}
+@model function observe_multivariate(
+    x = [1.5, 2.0, 2.5],
+    ::Type{TV} = Vector{Float64},
+) where {TV}
     a = TV(undef, length(x))
     a .~ Normal()
     x ~ MvNormal(a, I)
diff --git a/models/observe_submodel.jl b/models/observe_submodel.jl
@@ -1,7 +1,7 @@
 @model function inner2(x, a)
     x ~ Normal(a)
 end
-@model function observe_submodel(x=1.5)
+@model function observe_submodel(x = 1.5)
     a ~ Normal()
     _ignore ~ to_submodel(inner2(x, a))
 end
diff --git a/models/pdb_eight_schools_centered.jl b/models/pdb_eight_schools_centered.jl
@@ -4,9 +4,9 @@ sigma = [15, 10, 16, 11, 9, 11, 10, 18]
 
 @model function pdb_eight_schools_centered(J, y, sigma)
     mu ~ Normal(0, 5)
-    tau ~ truncated(Cauchy(0, 5); lower=0)
+    tau ~ truncated(Cauchy(0, 5); lower = 0)
     theta = Vector{Float64}(undef, J)
-    for i in 1:J
+    for i = 1:J
         theta[i] ~ Normal(mu, tau)
         y[i] ~ Normal(theta[i], sigma[i])
     end
diff --git a/models/pdb_eight_schools_noncentered.jl b/models/pdb_eight_schools_noncentered.jl
